"""基于当前大屏和会话历史流式回答问题"""

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 QA_SYSTEM_PROMPT


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

        parser = StrOutputParser()
        parts: list[str] = []
        async for chunk in get_llm().astream(
            [
                SystemMessage(content=QA_SYSTEM_PROMPT),
                HumanMessage(
                    content=state["input_context"].model_dump_json(exclude_unset=True)
                ),
            ]
        ):
            # 工具调用等无文本片段会得到空串
            if delta := parser.invoke(chunk):
                parts.append(delta)
                if writer is not None:
                    writer(
                        {
                            "type": "answer_delta",
                            "message_id": state.get("message_id"),
                            "delta": delta,
                        }
                    )
        answer = "".join(parts).strip()
        # 没有有效文本时报错，避免将空回答标记为处理完成
        if not answer:
            raise ValueError("问答模型未返回文本回答")

        updated_state["answer"] = answer
        updated_state["status"] = "completed"
        return updated_state