"""基于当前大屏和澄清历史回答问题"""

import json

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser
from langgraph.config import get_stream_writer

from app.clients.llm import get_llm
from app.processor.base import BaseNode
from app.processor.state import QueryState
from app.prompt.qa import OFFER_EXTRACT_PROMPT, QA_SYSTEM_PROMPT
from app.schemas.offer import ExtractedOffer, InformReply


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

        def emit(delta: str) -> None:
            if writer is not None:
                writer(
                    {
                        "type": "answer_delta",
                        "message_id": state.get("message_id"),
                        "delta": delta,
                    }
                )

        content = state["input_context"].model_dump_json(exclude_unset=True)
        offer = None
        # 请求无法执行或执行失败时，附上事实说明，由本节点如实告知用户
        if notice := state.get("notice"):
            content += f"\n\n系统说明：{notice}"
            # 回复中的替代修改需要结构化保存，下一轮用户同意时直接执行，因此不流式输出
            reply = await get_llm(InformReply).ainvoke(
                [SystemMessage(content=QA_SYSTEM_PROMPT), HumanMessage(content=content)]
            )
            answer, offer = reply.answer.strip(), reply.offer
            emit(answer)
        else:
            answer = await self._stream(content, emit)
            # 如果没有有效文本，抛出异常，避免将空回答标记为处理完成
            if not answer:
                raise ValueError("问答模型未返回文本回答")
            # 回答已流式输出，再提取其中等待用户确认的修改，下一轮同意时直接执行
            offer = await self._extract_offer(state["input_context"].prompt, answer)

        # 保存最终回答，供后续输出给用户
        updated_state["answer"] = answer
        # 保存等待用户确认的提议，没有提议时清空
        updated_state["offer"] = offer
        # 回答生成完成；失败收尾后的说明保持失败状态，不交付动作
        if state.get("status") != "failed":
            updated_state["status"] = "completed"

        # 返回包含回答和完成状态的更新结果
        return updated_state

    async def _extract_offer(self, prompt: str, answer: str) -> str | None:
        try:
            result = await get_llm(ExtractedOffer).ainvoke(
                [
                    SystemMessage(content=OFFER_EXTRACT_PROMPT),
                    HumanMessage(
                        content=json.dumps(
                            {"user": prompt, "assistant": answer}, ensure_ascii=False
                        )
                    ),
                ]
            )
        except Exception as error:
            # 回答已交付，提取失败只影响下一轮能否直接执行，不让本轮失败
            self.logger.warning(f"提取回答中的提议失败: {error!r}")
            return None
        return result.offer

    @staticmethod
    async def _stream(content: str, emit) -> str:
        # 获取普通对话模型，问答只需文本回答，不绑定结构化输出协议
        llm_client = get_llm()
        parser = StrOutputParser()
        parts: list[str] = []
        # 异步流式调用模型，结合用户问题、当前大屏和澄清记录生成回答
        async for chunk in llm_client.astream(
            [
                # 系统消息提供稳定的问答规则和回答边界
                SystemMessage(content=QA_SYSTEM_PROMPT),
                # 输入上下文序列化为 JSON，只保留显式设置的字段
                HumanMessage(content=content),
            ]
        ):
            # 从每个片段提取文本，兼容文本内容块，工具调用等无文本片段会得到空串
            delta = parser.invoke(chunk)
            if delta:
                parts.append(delta)
                emit(delta)
        # 拼接完整回答，并去除首尾的空白
        return "".join(parts).strip()