"""阶段 3 chat 端点: 接 LangGraph agent, 完整 AG-UI 9 类事件.

事件流顺序 (典型):
  thinking        {content: "用户在询问财务数据..."}
  agent_step      {node: "route", status: "running"}
  agent_step      {node: "route", status: "done"}
  retrieval       {doc_ids: [...], scores: [...]}
  citation        {doc_id, page, snippet, score}
  token           {content: "..."}      (多次, 流式)
  agent_step      {node: "evaluate", status: "done"}
  done            {usage, total_ms}
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from collections.abc import AsyncIterator

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from langchain_core.messages import HumanMessage

from app.agents.graph import get_compiled_graph
from app.agents.nodes import answer_node_stream
from app.agents.state import empty_state_for
from app.config import settings
from app.core.auth import require_access_token
from app.llm.base import LLMMessage
from app.llm.factory import get_llm
from app.models import db
from app.models.schemas import ChatRequest, ChatResponse
from app.services.persist import schedule_push
from app.streaming.events import (
    sse_done,
    sse_error,
    sse_format,
    EV_AGENT_STEP,
    EV_CITATION,
    EV_DONE,
    EV_ERROR,
    EV_PROGRESS,
    EV_RETRIEVAL,
    EV_THINKING,
    EV_TOKEN,
    EV_TOOL_CALL,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/chat", tags=["chat"], dependencies=[Depends(require_access_token)])

# 持引用防 GC: 后台算 title 的任务可能跑 1-2s, 中间被 GC 会丢更新
_deferred_tasks: set[asyncio.Task] = set()


# ========== 兜底: 阶段 1 简单直答 (route=direct / 无 agent) ==========
_FALLBACK_SYSTEM = (
    "你是一个有帮助的中文 AI 助手。"
    "回答应简洁、准确、礼貌。如不知道答案请直接说明。"
)


@router.post("", response_model=ChatResponse)
async def chat_simple(req: ChatRequest) -> ChatResponse:
    """非流式 chat: 跑完整 agent 但收集完整结果再返."""
    started = time.time()
    try:
        final_state, events_log = await _run_agent(req)
    except Exception as e:  # noqa: BLE001
        logger.exception("chat failed")
        raise HTTPException(status_code=502, detail=str(e)) from e

    answer = final_state.get("final_answer") or "(无回答)"
    return ChatResponse(
        session_id=req.session_id,
        message_id=uuid.uuid4().hex,
        content=answer,
        citations=final_state.get("citations", []),
        agent_steps=events_log,
        usage={},  # 阶段 3 暂不聚合
        total_ms=int((time.time() - started) * 1000),
    )


@router.post("/stream")
async def chat_stream(req: ChatRequest) -> StreamingResponse:
    """SSE 流式 chat: 完整 AG-UI 事件."""
    return StreamingResponse(
        _agent_event_stream(req),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "X-Accel-Buffering": "no",  # 关闭 nginx 缓冲
        },
    )


# ========== 核心: 跑 agent + 产生 SSE 事件 ==========
async def _agent_event_stream(req: ChatRequest) -> AsyncIterator[str]:
    """驱动 agent + 把过程映射为 AG-UI 事件."""
    queue: asyncio.Queue[tuple[str, dict] | None] = asyncio.Queue(maxsize=256)
    failed = False

    async def emit(event: str, data: dict) -> None:
        await queue.put((event, data))

    async def runner() -> None:
        nonlocal failed
        try:
            await _run_agent(req, emit=emit)
        except Exception as e:  # noqa: BLE001
            failed = True
            logger.exception("agent runner failed")
            await emit(EV_ERROR, {"code": "agent_failed", "message": str(e), "retryable": True})
        finally:
            await queue.put(None)

    task = asyncio.create_task(runner())

    # 启动心跳
    started = time.time()
    try:
        while True:
            try:
                item = await asyncio.wait_for(queue.get(), timeout=15.0)
            except asyncio.TimeoutError:
                # 心跳
                yield ": keepalive\n\n"
                continue
            if item is None:
                break
            event, data = item
            yield sse_format(event, data)
    finally:
        if not task.done():
            task.cancel()

    # `error` is terminal too.  Do not emit a contradictory done event after it.
    if not failed:
        yield sse_done(total_ms=int((time.time() - started) * 1000))


async def _run_agent(
    req: ChatRequest,
    emit=None,  # async callable | None
) -> tuple[dict, list[dict]]:
    """执行 agent. emit 可选, 用于 SSE 推送中间事件.

    Returns:
        (final_state_dict, events_log)
    """
    started = time.time()
    events_log: list[dict] = []

    async def _emit(event: str, data: dict) -> None:
        if emit is not None:
            await emit(event, data)
        # 非流式接口的 agent_steps 只保留真正的节点事件。旧实现把每个 token
        # 都塞进 events_log，长回答会产生数百条无意义的“步骤”并浪费内存。
        if event == EV_AGENT_STEP:
            events_log.append({"event": event, "data": data, "ts": time.time()})

    # 第一帧必须尽早发出：即使后续模型冷启动或检索较慢，用户也能立即看到反馈。
    await _emit(EV_PROGRESS, {"pct": 2, "label": "已收到问题，正在准备分析…"})

    # 1. 加载历史消息
    history_rows = db.message_list_by_session(req.session_id, limit=20)
    from app.agents.state import messages_to_lc
    history_msgs = messages_to_lc([
        {"role": r["role"], "content": r["content"]}
        for r in history_rows
    ])

    # 2. 添加本轮 user message
    new_user_msg = HumanMessage(content=req.message)
    history_msgs.append(new_user_msg)

    # 3. 初始化 state
    state = empty_state_for(
        req.session_id,
        locale=req.locale,
        requested_doc_ids=list(dict.fromkeys(req.doc_ids or [])),
    )
    state["messages"] = history_msgs

    # 4. 记录 user message 到 SQLite
    #    - user message 同步写 (审计需要, 不能丢)
    #    - session 已有 title → 直接 upsert
    #    - session 还没有 title → 用 asyncio.create_task 后台算, 不阻塞首字
    #      (注: 实际 _auto_title 是字符串截断, 不是 LLM, 收益微小; 但作为
    #       防御性编程, 把所有非关键路径工作都移出 hot path)
    existing = db.session_get(req.session_id)
    if existing is not None and existing.get("title"):
        db.session_upsert(req.session_id, title=None)  # 仅 updated_at 刷新
    else:
        # 首次消息: 先空 title 写一行占位, 后台异步算 title 后再 upsert 更新
        db.session_upsert(req.session_id, title=None)

        async def _deferred_title() -> None:
            try:
                from app.api.sessions import _auto_title
                title = _auto_title(req.message)
                db.session_upsert(req.session_id, title=title)
            except Exception as e:  # noqa: BLE001
                logger.warning("Deferred auto_title failed: %s", e)

        # 持引用防 GC
        task = asyncio.create_task(_deferred_title())
        _deferred_tasks.add(task)
        task.add_done_callback(_deferred_tasks.discard)
    db.message_insert({
        "id": uuid.uuid4().hex,
        "session_id": req.session_id,
        "role": "user",
        "content": req.message,
    })

    # 5. 跑 LangGraph (retrieval loop)
    try:
        await _emit(EV_PROGRESS, {"pct": 5, "label": "正在初始化 Agent…"})
        graph = await get_compiled_graph()
    except Exception as e:  # noqa: BLE001
        raise

    config = {"configurable": {"thread_id": req.session_id}}
    final_state: dict = dict(state)

    # 使用 astream(updates) 逐节点执行。旧实现 graph.ainvoke 会等 route → retrieve
    # → rerank → evaluate 全部结束才返回，导致前端在数秒到数十秒内完全无反馈。
    retrieval_emitted = False
    retrieval_round = 0
    await _emit(EV_AGENT_STEP, {"node": "route", "status": "running"})
    await _emit(EV_PROGRESS, {"pct": 8, "label": "正在理解问题与判断检索路径…"})
    try:
        async for update in graph.astream(state, config=config, stream_mode="updates"):
            if not isinstance(update, dict):
                continue
            for node_name, node_update in update.items():
                if not isinstance(node_update, dict):
                    continue
                final_state.update(node_update)
                await _emit(EV_AGENT_STEP, {"node": node_name, "status": "done"})

                if node_name == "route":
                    if final_state.get("route_decision") == "direct":
                        await _emit(EV_PROGRESS, {
                            "pct": 58,
                            "label": "问题无需知识库检索，正在组织回答…",
                        })
                    else:
                        await _emit(EV_AGENT_STEP, {"node": "retrieve", "status": "running"})
                        await _emit(EV_PROGRESS, {"pct": 18, "label": "正在检索知识库…"})

                elif node_name == "retrieve":
                    retrieval_round += 1
                    retrieved = final_state.get("retrieved", [])
                    doc_ids = final_state.get("retrieved_doc_ids", [])
                    await _emit(EV_RETRIEVAL, {
                        "doc_ids": doc_ids,
                        "count": len(retrieved),
                        "scores": [round(h.score, 4) for h in retrieved[:10]],
                    })
                    retrieval_emitted = True
                    await _emit(EV_AGENT_STEP, {"node": "rerank", "status": "running"})
                    await _emit(EV_PROGRESS, {
                        "pct": 38 if retrieval_round == 1 else 56,
                        "label": (
                            f"已找到 {len(retrieved)} 个候选片段，正在筛选证据…"
                            if retrieval_round == 1
                            else f"扩展检索命中 {len(retrieved)} 个片段，正在重新筛选…"
                        ),
                    })

                elif node_name == "rerank":
                    await _emit(EV_AGENT_STEP, {"node": "evaluate", "status": "running"})
                    await _emit(EV_PROGRESS, {
                        "pct": 52 if retrieval_round == 1 else 59,
                        "label": "正在校验资料相关性…",
                    })

                elif node_name == "evaluate":
                    if final_state.get("needs_more_retrieval"):
                        await _emit(EV_AGENT_STEP, {"node": "retrieve", "status": "running"})
                        await _emit(EV_PROGRESS, {
                            "pct": 54,
                            "label": "现有证据不足，正在扩展检索…",
                        })
                    else:
                        await _emit(EV_PROGRESS, {
                            "pct": 62,
                            "label": "证据校验完成，正在组织回答…",
                        })
    except Exception as e:  # noqa: BLE001
        logger.exception("graph.astream failed")
        raise

    # 6. 兼容性兜底: 正常 astream 已在 retrieve 节点完成时推送；若未来图实现
    # 不产生该 update，仍在生成答案前补发一次检索快照。
    if final_state.get("retrieved_doc_ids") and not retrieval_emitted:
        await _emit(EV_RETRIEVAL, {
            "doc_ids": final_state["retrieved_doc_ids"],
            "count": len(final_state.get("retrieved", [])),
            "scores": [round(h.score, 4) for h in final_state.get("retrieved", [])[:10]],
        })

    # 7. 跑 answer 流式 (直接调 LLM, 不走 graph)
    await _emit(EV_AGENT_STEP, {"node": "answer", "status": "running"})
    await _emit(EV_PROGRESS, {"pct": 68, "label": "正在生成回答，内容将实时显示…"})

    async def _on_token(content: str) -> None:
        await _emit(EV_TOKEN, {"content": content})

    async def _on_citation(c: dict) -> None:
        await _emit(EV_CITATION, c)

    async def _on_thinking(content: str) -> None:
        await _emit(EV_THINKING, {"content": content})

    answer_update = await answer_node_stream(
        final_state,
        on_token=_on_token,
        on_citation=_on_citation,
        on_thinking=_on_thinking,
    )
    final_state.update(answer_update)

    await _emit(EV_AGENT_STEP, {"node": "answer", "status": "done"})
    await _emit(EV_PROGRESS, {"pct": 100, "label": "完成"})

    # 8. 记录 assistant message 到 SQLite
    elapsed = int((time.time() - started) * 1000)
    db.message_insert({
        "id": uuid.uuid4().hex,
        "session_id": req.session_id,
        "role": "assistant",
        "content": final_state.get("final_answer", ""),
        "citations": final_state.get("citations", []),
    })

    # 9. ⚠️ 关键: 把 session/message 写完 → 调度 persist push (5s debounce).
    #    不调的话: 写本地 /data/sqlite/app.db → Space 重启 /data 临时卷清空 →
    #    历史对话全丢. 5s debounce 够合并同一 session 的 user+assistant 写入.
    try:
        await schedule_push()
    except Exception as e:  # noqa: BLE001
        logger.warning("schedule_push after chat failed: %s", e)

    return final_state, events_log