import uuid
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from api import services
from api.agent import Agent
from api.config import config


@asynccontextmanager
async def lifespan(_: FastAPI):
    await config.init_mongo_db()
    yield


app = FastAPI(lifespan=lifespan)
agent = Agent()


@app.post("/api/v1/conversation")
async def create_conversation():
    conversation_id = uuid.uuid4()
    await services.create_conversation(conversation_id, agent.system_prompt)
    return {"conversation_id": conversation_id}


class Message(BaseModel):
    conversation_id: uuid.UUID
    prompt: str

@app.post("/api/v1/message")
async def send(message: Message):
    def b64_sse(func):
        """Este helper transforma un generador de strings a un generador del protocolo SSE"""

        async def wrapper(*args, **kwargs):
            async for chunk in func(*args, **kwargs):
                content = chunk.model_dump_json()
                data = f"data: {content}\n\n"
                yield data

        return wrapper

    sse_stream = b64_sse(services.stream)
    generator = sse_stream(agent, message.prompt, message.conversation_id)
    return StreamingResponse(generator, media_type="text/event-stream")