"""POST /api/chat — SSE streaming chat with Agent."""

import json
import asyncio
import os
import traceback
from pathlib import Path
from typing import AsyncGenerator

from fastapi import APIRouter
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse

from graph.agent import agent_manager
from graph.session_manager import session_manager
from graph.chat_model_options import chat_model_options

router = APIRouter()


class ChatRequest(BaseModel):
    message: str
    session_id: str = "default"
    stream: bool = True


async def _generate_title(session_id: str) -> str | None:
    """Generate a title for a session using DeepSeek. Returns title or None."""
    try:
        messages = session_manager.load_session(session_id)
        first_user = ""
        first_assistant = ""
        for msg in messages:
            if msg["role"] == "user" and not first_user:
                first_user = msg["content"][:200]
            elif msg["role"] == "assistant" and not first_assistant:
                first_assistant = msg["content"][:200]
            if first_user and first_assistant:
                break

        if not first_user:
            return None

        from langchain_deepseek import ChatDeepSeek
        from langchain_core.messages import HumanMessage as HM

        llm = ChatDeepSeek(
            model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            api_base=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
            temperature=0.3,
            **chat_model_options(),
        )

        prompt = (
            f"根据以下对话内容，生成一个不超过10个字的中文标题，只输出标题文本，不要加引号或标点。\n\n"
            f"用户: {first_user}\n"
            f"助手: {first_assistant}"
        )

        result = await llm.ainvoke([HM(content=prompt)])
        title = result.content.strip().strip('"\'""''')[:20]
        session_manager.update_title(session_id, title)
        return title
    except Exception:
        traceback.print_exc()
        return None


class EmptyChatResponse(RuntimeError):
    pass


async def bounded_events(iterator, idle_seconds=75, total_seconds=180):
    """Bound stalls and total execution; never replay possible tool side effects."""
    deadline = asyncio.get_running_loop().time() + total_seconds
    try:
        while True:
            remaining = deadline - asyncio.get_running_loop().time()
            if remaining <= 0:
                raise TimeoutError()
            try:
                item = await asyncio.wait_for(anext(iterator), timeout=min(idle_seconds, remaining))
            except StopAsyncIteration:
                break
            yield item
    finally:
        await iterator.aclose()


async def event_generator(message: str, session_id: str) -> AsyncGenerator[dict, None]:
    """Generate SSE events from agent stream.

    Tracks multiple response segments — each time the agent finishes
    tool calls and starts generating new text, a new_response event is
    emitted and a new segment begins. Each segment is saved as a
    separate assistant message in the session history.
    """
    segments: list[dict] = []
    current_segment: dict = {"content": "", "tool_calls": []}
    saved = False
    completed = False
    try:
        # Use merged history for agent context (combines consecutive assistant msgs)
        history = session_manager.load_session_for_agent(session_id)
        is_first_message = len(history) == 0

        # Per-segment tracking
        segments: list[dict] = []  # finalized segments
        current_segment: dict = {"content": "", "tool_calls": []}

        async for event in bounded_events(agent_manager.astream(message, history, session_id)):
            event_type = event.get("type", "unknown")

            if event_type == "retrieval":
                yield {
                    "event": "retrieval",
                    "data": json.dumps(
                        {"query": event["query"], "results": event["results"]},
                        ensure_ascii=False,
                    ),
                }

            elif event_type == "token":
                current_segment["content"] += event["content"]
                yield {
                    "event": "token",
                    "data": json.dumps({"content": event["content"]}, ensure_ascii=False),
                }

            elif event_type == "new_response":
                # Finalize current segment and start a new one
                segments.append(current_segment)
                current_segment = {"content": "", "tool_calls": []}
                yield {
                    "event": "new_response",
                    "data": json.dumps({}, ensure_ascii=False),
                }

            elif event_type == "tool_start":
                current_segment["tool_calls"].append({
                    "tool": event["tool"],
                    "input": event.get("input", ""),
                })
                yield {
                    "event": "tool_start",
                    "data": json.dumps(
                        {"tool": event["tool"], "input": event["input"]},
                        ensure_ascii=False,
                    ),
                }

            elif event_type == "tool_end":
                # Match with current segment's tool_calls to add output
                for tc in reversed(current_segment["tool_calls"]):
                    if tc["tool"] == event["tool"] and "output" not in tc:
                        tc["output"] = event["output"]
                        break
                yield {
                    "event": "tool_end",
                    "data": json.dumps(
                        {"tool": event["tool"], "output": event["output"]},
                        ensure_ascii=False,
                    ),
                }

            elif event_type == "done":
                if not str(event.get("content") or "").strip():
                    raise EmptyChatResponse()
                # Some providers deliver final content without token events.
                if not current_segment["content"] and not any(s["content"] for s in segments):
                    current_segment["content"] = event["content"]
                    yield {"event":"token", "data":json.dumps({"content":event["content"]}, ensure_ascii=False)}
                # Finalize the last segment
                segments.append(current_segment)

                # Save user message once
                session_manager.save_message(session_id, "user", message)

                # Save each segment as a separate assistant message
                for seg in segments:
                    if not seg["content"].strip() and not seg["tool_calls"]:
                        continue
                    tc = seg["tool_calls"] if seg["tool_calls"] else None
                    session_manager.save_message(
                        session_id, "assistant", seg["content"], tool_calls=tc
                    )
                saved = True
                completed = True

                yield {
                    "event": "done",
                    "data": json.dumps(
                        {"content": event["content"], "session_id": session_id},
                        ensure_ascii=False,
                    ),
                }

                # Auto-generate title on first conversation
                if is_first_message:
                    try:
                        title = await asyncio.wait_for(_generate_title(session_id), timeout=8)
                    except TimeoutError:
                        title = None
                    if title:
                        yield {
                            "event": "title",
                            "data": json.dumps(
                                {"session_id": session_id, "title": title},
                                ensure_ascii=False,
                            ),
                        }

        if not completed:
            raise ConnectionError("stream ended before done")
    except Exception as e:
        is_timeout = isinstance(e, TimeoutError) or type(e).__name__ in ("APITimeoutError", "ReadTimeout", "ConnectTimeout")
        code = "empty_response" if isinstance(e, EmptyChatResponse) else "timeout" if is_timeout else "chat_failed"
        message_text = {
            "empty_response":"模型未返回回答正文，可能只返回了推理字段。请检查聊天模型兼容配置后重试。",
            "timeout":"模型响应超时，已停止本次等待。请检查服务后重试；不会自动重复执行工具。",
            "chat_failed":"聊天未完成，请检查当前后端、模型连接或访问凭据后重试。已接收的内容会保留。",
        }[code]
        print(f"Chat failed: {type(e).__name__}")
        if not saved:
            session_manager.save_message(session_id, "user", message)
            for seg in segments:
                if seg["content"] or seg["tool_calls"]:
                    session_manager.save_message(session_id, "assistant", seg["content"], tool_calls=seg["tool_calls"] or None)
            session_manager.save_message(session_id, "assistant", current_segment["content"] + "\n\n【未完成】" + message_text, tool_calls=current_segment["tool_calls"] or None)
        yield {
            "event": "error",
            "data": json.dumps({"error": message_text, "code": code}, ensure_ascii=False),
        }


@router.post("/chat")
async def chat(request: ChatRequest):
    if request.stream:
        return EventSourceResponse(
            event_generator(request.message, request.session_id)
        )
    # Non-streaming fallback
    result = await agent_manager.ainvoke(request.message, request.session_id)
    return {"reply": result}


@router.get("/chat/raw-traces")
async def get_raw_traces(session_id: str = "default", limit: int = 20):
    """获取指定 session 的原始 LLM 调用报文

    Args:
        session_id: 会话 ID
        limit: 返回最近 N 条记录 (默认 20)

    Returns:
        {"traces": [...]} 原始报文列表，按时间倒序
    """
    try:
        from graph.raw_trace import load_raw_traces

        # 从 agent_manager 获取 raw_logs_dir
        raw_logs_dir = agent_manager._raw_logs_dir
        if not raw_logs_dir:
            return {"traces": []}

        traces = load_raw_traces(session_id, raw_logs_dir, limit=limit)
        return {"traces": traces}
    except Exception as e:
        traceback.print_exc()
        return {"traces": [], "error": str(e)}


@router.get("/chat/trajectory")
async def get_trajectory(session_id: str = "default"):
    """获取会话完整交互轨迹 - 包含完整的 human/assistant/tool 消息流

    Args:
        session_id: 会话 ID

    Returns:
        {
          "session_id": "...",
          "messages": [
            {"role": "human", "content": "...", "timestamp": ...},
            {"role": "assistant", "content": "...", "timestamp": ..., "tool_calls": [...]},
            {"role": "tool", "name": "...", "content": "...", "timestamp": ...},
            {"role": "assistant", "content": "...", "timestamp": ...}
          ],
          "steps": [...]  // 保留旧格式用于向后兼容
        }
    """
    try:
        import json
        from graph.raw_trace import load_raw_traces

        raw_logs_dir = agent_manager._raw_logs_dir
        if not raw_logs_dir:
            # 无 raw_logs_dir，返回空
            return {"session_id": session_id, "messages": [], "steps": []}

        # 优先读取完整消息流日志文件
        msgs_log_file = raw_logs_dir / f"{session_id}.msgs.jsonl"
        messages = []

        if msgs_log_file.exists():
            # 从 msgs.jsonl 读取完整消息流（包含 tool 消息）
            with open(msgs_log_file, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            msg = json.loads(line)
                            messages.append(msg)
                        except json.JSONDecodeError:
                            continue
        else:
            # 回退到旧逻辑：从 session 文件读取（无 tool 消息）
            session_messages = session_manager.load_session(session_id)
            for msg in session_messages:
                messages.append({
                    "role": msg["role"],
                    "content": msg["content"],
                    "timestamp": None,
                })

        # 构建旧格式 steps（向后兼容）
        steps = []
        traces = load_raw_traces(session_id, raw_logs_dir, limit=1000)
        traces.sort(key=lambda x: x.get("timestamp_start", 0))

        # 按 user/assistant 消息构建 steps（忽略 tool 消息）
        session_messages = session_manager.load_session(session_id)
        assistant_count = sum(1 for m in session_messages if m["role"] == "assistant")

        for msg_idx, msg in enumerate(session_messages):
            step = {
                "role": msg["role"],
                "content": msg["content"],
                "timestamp": None,
            }

            if msg["role"] == "assistant":
                llm_calls = []
                assistant_seq = sum(1 for m in session_messages[:msg_idx+1] if m["role"] == "assistant")
                traces_per_assistant = len(traces) / max(assistant_count, 1)
                start_trace = int((assistant_seq - 1) * traces_per_assistant)
                end_trace = int(assistant_seq * traces_per_assistant)

                for i in range(start_trace, min(end_trace, len(traces))):
                    trace = traces[i]
                    llm_call = {
                        "model": trace.get("request", {}).get("model", "unknown"),
                        "finish_reason": trace.get("finish_reason", "unknown"),
                        "request_messages": trace.get("request", {}).get("messages", []),
                        "response_text": "",
                        "tool_calls_list": [],
                        "usage": trace.get("usage", {}),
                        "duration_ms": trace.get("duration_ms", 0),
                    }

                    response_gens = trace.get("response", {}).get("generations", [])
                    if response_gens:
                        llm_call["response_text"] = response_gens[0].get("text", "")
                        msg_data = response_gens[0].get("message", {})
                        if isinstance(msg_data, dict) and "tool_calls" in msg_data:
                            tool_calls = msg_data["tool_calls"]
                            if tool_calls:
                                llm_call["tool_calls_list"] = [
                                    {
                                        "name": tc.get("name", "") if isinstance(tc, dict) else getattr(tc, "name", ""),
                                        "args": tc.get("args", {}) if isinstance(tc, dict) else getattr(tc, "args", {}),
                                    }
                                    for tc in tool_calls
                                ]

                    llm_calls.append(llm_call)
                    if step["timestamp"] is None:
                        step["timestamp"] = trace.get("timestamp_start")

                step["llm_calls"] = llm_calls

            steps.append(step)

        return {
            "session_id": session_id,
            "messages": messages,  # 新格式：完整消息流（含 tool）
            "steps": steps,  # 旧格式：向后兼容
        }

    except Exception as e:
        traceback.print_exc()
        return {"session_id": session_id, "messages": [], "steps": [], "error": str(e)}