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

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

import json

from langchain.agents.middleware.model_call_limit import ModelCallLimitExceededError
from langchain_core.messages import AIMessageChunk, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langgraph.config import get_stream_writer

from app.agent.data_tools import DataContext
from app.agent.qa_agent import QAContext, get_qa_agent
from app.processor.base import BaseNode
from app.processor.history import history_messages
from app.processor.state import QueryState
from app.processor.summary import data_file_summary

# 模型在一次回答里分几次输出文字（例如查询前后各说一段）时，段落之间空一行
PARAGRAPH_BREAK = "\n\n"


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

        parser = StrOutputParser()
        parts: list[str] = []
        last_message_id = None
        try:
            async for chunk, metadata in agent.astream(
                {
                    "messages": [
                        # 会话历史作为此前的对话消息，本次请求是最后一条用户消息
                        *history_messages(context.history),
                        HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
                    ]
                },
                context=QAContext(data=DataContext(tuple(context.data_files))),
                stream_mode="messages",
            ):
                # 工具结果、工具调用和中间件插入的消息都不推送
                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"
        return updated_state