"""SSE streaming chat endpoint."""

from __future__ import annotations

import json
import uuid
from collections.abc import AsyncIterator

from langchain_core.messages import HumanMessage
from litestar import post
from litestar.response import Stream

from app.graph.builder import get_chat_graph
from app.schemas import ChatRequest
from app.store.postgres import Conversation, Message, async_session_factory
from app.tools import resolve_tools
from app.tools.search import set_search_model


async def _ensure_conversation(conversation_id: str, first_message: str) -> None:
    """Create conversation and persist the user message."""
    async with async_session_factory() as session:
        existing = await session.get(Conversation, conversation_id)
        if existing is None:
            # Use first ~50 chars of message as title
            title = first_message[:50].strip() or "New conversation"
            conv = Conversation(id=conversation_id, title=title)
            session.add(conv)
        # Persist user message
        msg = Message(
            conversation_id=conversation_id,
            role="human",
            content=first_message,
        )
        session.add(msg)
        await session.commit()


async def _persist_ai_message(conversation_id: str, content: str) -> None:
    """Persist the AI response message."""
    async with async_session_factory() as session:
        msg = Message(
            conversation_id=conversation_id,
            role="ai",
            content=content,
        )
        session.add(msg)
        await session.commit()


async def _stream_response(request: ChatRequest) -> AsyncIterator[bytes]:
    """Stream LLM response tokens via SSE."""
    # Ensure conversation exists and persist user message
    await _ensure_conversation(request.conversation_id, request.message)

    # Resolve tools from frontend tool keys
    active_tools = resolve_tools(request.tools)

    # Set search strategy model so the web_search tool knows depth
    set_search_model(request.model)

    graph = await get_chat_graph(model=request.model, tools=active_tools)

    config = {
        "configurable": {"thread_id": request.conversation_id},
    }

    # When using ReAct agent (with tools), input is just messages.
    # When using plain graph (no tools), input includes model key.
    if active_tools:
        input_data = {"messages": [HumanMessage(content=request.message)]}
    else:
        input_data = {
            "messages": [HumanMessage(content=request.message)],
            "model": request.model,
        }

    full_content: list[str] = []

    async for event in graph.astream_events(
        input_data,
        config=config,
        version="v2",
    ):
        kind = event.get("event", "")

        if kind == "on_chat_model_stream":
            chunk = event.get("data", {}).get("chunk")
            if chunk and hasattr(chunk, "content") and chunk.content:
                # Only stream text content, skip tool call chunks
                if isinstance(chunk.content, str):
                    full_content.append(chunk.content)
                    sse_data = json.dumps(
                        {"type": "token", "content": chunk.content},
                        ensure_ascii=False,
                    )
                    yield f"data: {sse_data}\n\n".encode("utf-8")

        elif kind == "on_tool_start":
            # Notify frontend that a tool is being called
            tool_name = event.get("name", "unknown")
            sse_data = json.dumps(
                {"type": "tool_start", "tool": tool_name},
                ensure_ascii=False,
            )
            yield f"data: {sse_data}\n\n".encode("utf-8")

        elif kind == "on_tool_end":
            tool_name = event.get("name", "unknown")
            sse_data = json.dumps(
                {"type": "tool_end", "tool": tool_name},
                ensure_ascii=False,
            )
            yield f"data: {sse_data}\n\n".encode("utf-8")

    # Persist AI response
    ai_content = "".join(full_content)
    if ai_content:
        await _persist_ai_message(request.conversation_id, ai_content)

    # Send done signal
    done_data = json.dumps({"type": "done"})
    yield f"data: {done_data}\n\n".encode("utf-8")


@post("/api/chat/stream")
async def stream_chat(data: ChatRequest) -> Stream:
    """POST /api/chat/stream - SSE streaming chat endpoint."""
    if not data.conversation_id:
        data.conversation_id = str(uuid.uuid4())

    return Stream(
        _stream_response(data),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )