"""基于当前大屏、上传的数据和会话历史流式回答问题

由问答 agent 执行：会话里有数据文件时先查询再回答，否则只调用一次模型；只把模型输出的文字推送给前端
"""

import json

from langchain.agents.middleware.model_call_limit import ModelCallLimitExceededError
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
    ToolMessage,
)
from langchain_core.output_parsers import StrOutputParser
from langgraph.config import get_stream_writer

from app.agent.qa_agent import QAContext, get_qa_agent
from app.processor.base import BaseNode
from app.processor.history import history_messages, history_records
from app.processor.state import QueryState, run_data
from app.processor.summary import data_file_summary
from app.schemas.base import JsonObject
from app.schemas.conversation import TurnQuery

# 模型在一次回答里分几次输出文字（例如查询前后各说一段）时，段落之间空一行
PARAGRAPH_BREAK = "\n\n"
# 带数据的问答回答完后给的后续请求：意图判断会把"展示一下 9 月销售额"这类说法分到问答，
# 用户其实想在大屏上看到；回答之后给一个按钮，点了按普通请求走大屏 agent，会话历史里有这次回答
PUT_ON_SCREEN = "把上面的结果放到大屏上"
# 一轮最多记下的查询；Langfuse 里每条 trace 的 query_data 调用 p50 3 次、p90 7 次（含重试）
MAX_QUERY_RECORDS = 10


class QANode(BaseNode):
    name = "qa_node"

    async def process(self, state: QueryState) -> QueryState:
        updated_state: QueryState = state.copy()

        try:
            # 图执行上下文中通过 custom 流逐段推送回答，独立调用节点时没有写入器
            writer = get_stream_writer()
        except RuntimeError:
            writer = None

        context = state["input_context"]
        payload = context.model_dump(
            mode="json", exclude_unset=True, exclude={"history", "data_files"}
        )
        if context.data_files:
            # 只给表名、行数和列名，文件路径和列概况不进提示词，需要时由 describe_data 查看
            payload["data_files"] = [
                data_file_summary(item) for item in context.data_files
            ]
        agent = get_qa_agent(with_data=bool(context.data_files))

        # 此前各轮的修改、查询和澄清记录跟着本次请求，不拼进历史消息；前端带来的同名字段不保留
        payload.pop("history_records", None)
        if records := history_records(context.history, context.screen):
            payload["history_records"] = records

        parser = StrOutputParser()
        parts: list[str] = []
        last_message_id = None
        messages: list[BaseMessage] = []
        try:
            async for mode, data in agent.astream(
                {
                    "messages": [
                        # 会话历史作为此前的对话消息，本次请求是最后一条用户消息
                        *history_messages(context.history),
                        HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
                    ]
                },
                context=QAContext(data=run_data()),
                # messages 推送文字，values 留下完整的消息，用来记下执行过的查询
                stream_mode=["messages", "values"],
            ):
                if mode == "values":
                    messages = data["messages"]
                    continue
                chunk, metadata = data
                # 工具结果、工具调用和中间件插入的消息都不推送
                if metadata.get("langgraph_node") != "model" or not isinstance(
                    chunk, AIMessageChunk
                ):
                    continue
                if not (delta := parser.invoke(chunk)):
                    continue
                if parts and chunk.id != last_message_id:
                    delta = PARAGRAPH_BREAK + delta
                last_message_id = chunk.id
                parts.append(delta)
                if writer is not None:
                    writer(
                        {
                            "type": "answer_delta",
                            "message_id": state.get("message_id"),
                            "delta": delta,
                        }
                    )
        except ModelCallLimitExceededError:
            self.logger.warning("问答 agent 超出模型调用上限（MAX_QA_STEPS）")
            updated_state["status"] = "failed"
            updated_state["answer"] = (
                "查询步骤超出上限，这次没有完成，可以缩小范围或换个问法再试"
            )
            return updated_state

        # 推送了什么就保存什么，会话历史和前端看到的回答一致
        answer = "".join(parts).strip()
        # 没有有效文本时报错，避免将空回答标记为处理完成
        if not answer:
            raise ValueError("问答模型未返回文本回答")

        updated_state["answer"] = answer
        updated_state["status"] = "completed"
        if context.data_files:
            updated_state["suggestions"] = [PUT_ON_SCREEN]
            # 只记成功的问答：步数用尽时回答里也没有数字可用
            updated_state["queries"] = query_records(messages)
        return updated_state


def query_records(messages: list[BaseMessage]) -> list[JsonObject]:
    """问答里成功执行的 query_data：调用参数和工具结果按 tool_call_id 配对

    只存 SQL 和结果的形状；按执行顺序，相同的 SQL 只留最后一次，最多 MAX_QUERY_RECORDS 条（留最近的）
    """
    sql_of = {
        call["id"]: call["args"].get("sql")
        for message in messages
        if isinstance(message, AIMessage)
        for call in message.tool_calls
        if call["name"] == "query_data"
    }
    found: dict[str, JsonObject] = {}
    for message in messages:
        if not isinstance(message, ToolMessage) or message.status == "error":
            continue
        sql = sql_of.get(message.tool_call_id)
        if not isinstance(sql, str):
            continue
        try:
            result = json.loads(str(message.content))
            record = TurnQuery(
                sql=sql,
                columns=result["columns"],
                rows=result["row_count"],
                truncated=result["truncated"],
            ).model_dump(mode="json")
        except (ValueError, KeyError, TypeError):
            continue
        # 重新插入，相同 SQL 的位置按最后一次执行算；按原文比较，字面量里的空白也有意义
        found.pop(sql, None)
        found[sql] = record
    return list(found.values())[-MAX_QUERY_RECORDS:]