import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

import logfire
from fastapi import FastAPI, Response
from nats.js import JetStreamContext

from examples.customer_support.customer_support import (
    CustomerSupportRequest,
    run_customer_support_agent,
)
from tfy_assistant_framework.assistant import (
    AssistantServe,
    CreateAssistantTaskRunResult,
    NATSAssistantUpdateSink,
    PostAssistantTaskMessage,
)
from tfy_assistant_framework.nats_client import nats_connection

js: JetStreamContext
assistant_serve = AssistantServe()


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    global js

    async with nats_connection() as (nc_, js):
        js = js

        async with assistant_serve.init(nc=nc_, app=app):
            yield


app = FastAPI(
    title="LLM Assistants Server",
    lifespan=lifespan,
    docs_url="/",
)


@app.post("/assistants/customer_support/task")
async def customer_support_task(
    body: CustomerSupportRequest,
) -> CreateAssistantTaskRunResult:
    task_id = uuid.uuid4().hex
    assistant_output = NATSAssistantUpdateSink(task_id=task_id, js=js)
    with logfire.span(
        "Customer Support Assistant",
        task_id=task_id,
    ):
        return await assistant_serve.create_assistant_task(
            task_id=task_id,
            assistant_awaitable=run_customer_support_agent(body),
            assistant_update_sink=assistant_output,
        )


@app.post("/tasks/{task_id}/message")
async def send_message_to_assistant_task(
    task_id: str,
    message: PostAssistantTaskMessage,
) -> Response:
    return await assistant_serve.send_message_to_task(task_id, message)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("examples.customer_support.app:app", host="0.0.0.0", port=8000)
