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

import json
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

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,
        )

        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


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.
    """
    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 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":
                # 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:
                    tc = seg["tool_calls"] if seg["tool_calls"] else None
                    session_manager.save_message(
                        session_id, "assistant", seg["content"], tool_calls=tc
                    )

                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:
                    title = await _generate_title(session_id)
                    if title:
                        yield {
                            "event": "title",
                            "data": json.dumps(
                                {"session_id": session_id, "title": title},
                                ensure_ascii=False,
                            ),
                        }

    except Exception as e:
        traceback.print_exc()
        yield {
            "event": "error",
            "data": json.dumps({"error": str(e)}, 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)}