import asyncio
from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
from typing import Optional
from PyCharacterAI import get_client
from PyCharacterAI.exceptions import SessionClosedError, RequestError
import uvicorn
import os

app = FastAPI()


DEFAULT_TOKEN = os.getenv("token")
DEFAULT_CHARACTER_ID = "smtV3Vyez6ODkwS8BErmBAdgGNj-1XWU73wIFVOY1hQ"
DEFAULT_VOICE_ID = "26edcb78-7e4d-4d81-970e-6fdaf81ed778"  # Replace with the desired voice ID

# This will hold the client instance, which should be initialized once
client = None

@app.on_event("startup")
async def startup_event():
    global client
    # Initialize client during startup
    client = await get_client(token=DEFAULT_TOKEN)

@app.on_event("shutdown")
async def shutdown_event():
    # Close the client session during shutdown
    await client.close_session()

async def initialize_client(token: str):
    """Helper function to reinitialize the client connection."""
    global client
    if client is not None:
        await client.close_session()  # Ensure any existing session is closed
    client = await get_client(token=token)

@app.get("/send_message")
async def send_message(
    message: str,
    token: Optional[str] = Query(None, description="API token for authentication."),
    character_id: str = Query(DEFAULT_CHARACTER_ID, description="Character ID for the chat session."),
    chat_id: Optional[str] = Query(None, description="ID of the existing chat session, if available."),
    voice_id: str = Query(DEFAULT_VOICE_ID, description="Voice ID for generating speech."),
    voice: bool = Query(True, description="Set to true to generate voice, false to skip voice generation.")
):
    """
    Send a message to the character. If no chat_id is provided, initialize a new chat session.
    Optionally generate voice for the response.
    """
    try:
  
        token = token or DEFAULT_TOKEN

        # Ensure client connection is active
        if client is None:
            await initialize_client(token)

        # If chat_id is not provided, create a new chat session
        if not chat_id:
            chat, greeting_message = await client.chat.create_chat(character_id)
            chat_id = chat.chat_id
            greeting_text = {
                "author": greeting_message.author_name,
                "text": greeting_message.get_primary_candidate().text
            }
        else:
            greeting_text = None  # No greeting if we’re using an existing chat

        # Send the message and get the response
        answer = await client.chat.send_message(character_id, chat_id, message)
        response_text = answer.get_primary_candidate().text

        speech_url = None  # Default to None
        if voice:  # Only generate voice if the voice parameter is true
            try:
                speech_url = await client.utils.generate_speech(
                    chat_id,
                    answer.turn_id,
                    answer.get_primary_candidate().candidate_id,
                    voice_id,
                    return_url=True
                )
            except Exception as e:
                # Log the voice generation error and continue without the speech URL
                print(f"Voice generation failed: {e}")

        # Prepare the response data
        response_data = {
            "chat_id": chat_id,
            "author": answer.author_name,
            "response": response_text,
            "voice_url": speech_url if voice else None
        }

        # Include greeting message only if a new chat was created
        if greeting_text:
            response_data["greeting_message"] = greeting_text

        return response_data

    except (SessionClosedError, RequestError):
        # Attempt to reinitialize the client on connection closure
        print("Connection closed. Reinitializing client.")
        await initialize_client(token)
        raise HTTPException(status_code=500, detail="Connection was closed. Please try again.")

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Root endpoint
@app.get("/")
def root():
    return {"message": "Welcome to the Character AI API with FastAPI!"}

# Run with: uvicorn main:app --host 0.0.0.0 --port 8000
if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860, workers=8, timeout_keep_alive=60000)