import json
import uuid
from collections.abc import AsyncIterator

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI

app = FastAPI()
model = ChatOpenAI(model="gpt-4o-mini", temperature=0)


def to_sse(part: dict) -> str:
    return f"data: {json.dumps(part, ensure_ascii=False)}\n\n"


def get_last_user_text(messages: list[dict]) -> str:
    for message in reversed(messages):
        if message.get("role") != "user":
            continue

        parts = message.get("parts", [])
        for part in parts:
            if part.get("type") == "text":
                return part.get("text", "")

        content = message.get("content")
        if isinstance(content, str):
            return content

    return ""


def chunk_to_text(chunk) -> str:
    content = getattr(chunk, "content", "")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            item.get("text", "")
            for item in content
            if isinstance(item, dict) and item.get("type") == "text"
        )
    return ""


@app.post("/chat")
async def chat(request: Request) -> StreamingResponse:
    body = await request.json()
    prompt = get_last_user_text(body.get("messages", []))

    async def event_stream() -> AsyncIterator[str]:
        message_id = f"msg_{uuid.uuid4().hex}"
        text_id = f"text_{uuid.uuid4().hex}"

        yield to_sse({"type": "start", "messageId": message_id})
        yield to_sse({"type": "start-step"})
        yield to_sse({
            "type": "data-status",
            "data": {
                "stage": "langchain",
                "message": "LangChain 已开始执行模型调用",
            },
        })
        yield to_sse({"type": "text-start", "id": text_id})

        async for chunk in model.astream([HumanMessage(content=prompt)]):
            delta = chunk_to_text(chunk)
            if delta:
                yield to_sse({
                    "type": "text-delta",
                    "id": text_id,
                    "delta": delta,
                })

        yield to_sse({"type": "text-end", "id": text_id})
        yield to_sse({
            "type": "data-status",
            "data": {
                "stage": "done",
                "message": "本轮回答已完成",
            },
        })
        yield to_sse({"type": "finish-step"})
        yield to_sse({"type": "finish"})

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "x-vercel-ai-ui-message-stream": "v1",
        },
    )

from langgraph_agent_kit import Orchestrator, OrchestratorHooks, StreamEndInfo

# 1. 实现 AgentRunner
class MyAgentRunner:
    async def run(self, message, context, **kwargs):
        """通过 context.emitter 发送事件"""
        agent = build_my_agent(context.emitter)
        async for chunk in agent.astream({"messages": [{"role": "user", "content": message}]}):
            await handler.handle(chunk)
        await context.emitter.aemit("__end__", None)

# 2. 定义钩子
async def save_to_db(info: StreamEndInfo):
    await db.save_message(
        conversation_id=info.conversation_id,
        content=info.aggregator.full_content,
        tool_calls=info.aggregator.tool_calls_list,
    )

# 3. 创建编排器
orchestrator = Orchestrator(
    agent_runner=MyAgentRunner(),
    hooks=OrchestratorHooks(on_stream_end=save_to_db),
)

# 4. 在路由中使用
async for event in orchestrator.run(message="你好", conversation_id="c1", user_id="u1"):
    yield event