import uuid
import os

# fastapi
from fastapi import APIRouter, Body, HTTPException, Header, Depends, Request
from pydantic import BaseModel
from fastapi.responses import StreamingResponse
from gai.lib.common.errors import InternalException
from gai.lib.common.logging import getLogger
logger = getLogger(__name__)
from gai.persona.profile.pydantic.ProvisionAgentPydantic import ProvisionAgentPydantic
from gai.persona.profile.pydantic.GenerateDescriptionPydantic import GenerateDescriptionPydantic
from gai.persona.persona_builder import PersonaBuilder
from gai.lib.common.utils import get_app_path
from gai.ttt.client.ttt_client import TTTClient

# Implementations Below
profile_router = APIRouter()

### POST /api/v1/persona/provision
@profile_router.post("/api/v1/persona/provision")
async def post_persona_provision(req: ProvisionAgentPydantic=Body(...)):
    try:
        persona_builder = PersonaBuilder(
            provision=req
        )
        persona=persona_builder.build()
        persona.agent_profile.CustomPrompt
        app_path = get_app_path()
        export_dir=os.path.abspath(os.path.join(app_path,"persona",persona.agent_profile.Id))
        await persona_builder.export_async(export_dir=export_dir)
        return {"message": "Agent provisioned successfully.","agent_id":persona.agent_profile.Id}
    except Exception as e:
        id = str(uuid.uuid4())
        logger.error(str(e)+" "+id)
        raise InternalException(id)

### POST /api/v1/persona/description
@profile_router.post("/api/v1/persona/description")
async def post_persona_description(req: GenerateDescriptionPydantic=Body(...)):
    try:
        ttt_client=TTTClient({
            "type": "ttt",
            "url": "http://localhost:12031/gen/v1/chat/completions",
            "timeout": 60.0,
            "max_new_tokens": 500,
            "stop_conditions": ["<s>", "</s>", "user:", "\n\n"]
        })
        response = ttt_client(messages=
            f"""system: You are a role-play character generator and you are tasked to create a self-introduction for a fictitious AI character based on the name and personality traits provided by the user. You should assimilate the traits into the self-intro and stay in character.
             user: Generate persona with name {req.Name} and traits{str(req.AgentTraits)}.
             assistant:"""
        ,stream=True)
        def streamer():
            try:
                for chunk in response:
                    chunk=chunk.extract()
                    if type(chunk) is str:
                        yield chunk
            except Exception as e:
                logger.error("Streaming error: "+str(e))
                yield "Error during streaming"
        return StreamingResponse(streamer())
    except Exception as e:
        id = str(uuid.uuid4())
        logger.error(str(e)+" "+id)
        raise InternalException(id)

### GET /api/v1/persona/profile
@profile_router.get("/api/v1/persona/profile/{persona_id}")
async def get_persona_profile(persona_id:str):
    builder = PersonaBuilder()
    app_path = get_app_path()
    import_dir=os.path.abspath(os.path.join(app_path,"persona",persona_id))
    await builder.import_async(import_dir=import_dir)
    persona=builder.build()
    return persona.agent_profile
