from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from hugchat import hugchat
from hugchat.login import Login
import os

app = FastAPI()

# Pydantic model for request body
class QuestionRequest(BaseModel):
    question: str

# Global variable to store the chatbot instance
chatbot = None

def setup_chatbot(email, password, cookie_path, assistant_id):
    """
    Sets up the Hugging Face chatbot with login and conversation.
    
    Args:
        email (str): User email for login
        password (str): User password for login
        cookie_path (str): Directory to store cookies
        assistant_id (str): ID of the assistant to use
    
    Returns:
        hugchat.ChatBot: Configured chatbot instance
    """
    try:
        # Create cookie directory if it doesn't exist
        os.makedirs(cookie_path, exist_ok=True)
        
        sign = Login(email, password)
        cookies = sign.login(cookie_dir_path=cookie_path, save_cookies=True)
        chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
        chatbot.new_conversation(assistant=assistant_id, switch_to=True)
        return chatbot
    except Exception as e:
        raise Exception(f"Failed to set up chatbot: {e}")

# Initialize chatbot at startup
@app.on_event("startup")
async def startup_event():
    global chatbot
    # Credentials and configuration
    EMAIL = os.getenv("EMAIL")
    PASSWD = os.getenv("PASSWD")
    COOKIE_PATH_DIR = "./cookies/"
    ASSISTANT_ID = "682e0c1f5f0c3d952a27498e"  # Replace with your actual assistant ID
    
    chatbot = setup_chatbot(EMAIL, PASSWD, COOKIE_PATH_DIR, ASSISTANT_ID)

@app.post("/generate")
async def generate_response(request: QuestionRequest):
    """
    Generates a response from the AI based on the provided question.
    
    Args:
        request (QuestionRequest): JSON body containing the question.
    
    Returns:
        dict: A dictionary containing the AI's response or an error message.
    """
    global chatbot
    if chatbot is None:
        raise HTTPException(status_code=500, detail="Chatbot not initialized. Please try again later.")
    
    try:
        # Generate response (non-streaming for simplicity)
        response_data = chatbot.chat(request.question, stream=False)
        
        # Extract the actual response text
        # The response may be a dictionary; check for 'gen' or other keys
        if isinstance(response_data, dict):
            # Assuming 'gen' contains the response text (list of strings)
            response_text = "".join(response_data.get("gen", [])) if response_data.get("gen") else "Here's what we can do: Let's discuss your vision!"
            if not response_text:
                # Fallback to a default premium response if 'gen' is empty
                response_text = f"Welcome, valued client! How can Abdullah Ali and our premium team bring your vision to life with a custom website or AI chatbot?"
        else:
            response_text = response_data  # Direct string response (if hugchat returns string)
        
        return {"response": response_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to generate response: {str(e)}")