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

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser

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()
        llm_client = get_llm()
        llm_res = await llm_client.ainvoke(
            [
                SystemMessage(content=QA_SYSTEM_PROMPT),
                HumanMessage(
                    content=state["input_context"].model_dump_json(exclude_unset=True)
                ),
            ]
        )
        # 普通文本输出，同时兼容模型返回的文本内容块
        answer = StrOutputParser().invoke(llm_res).strip()
        if not answer:
            raise ValueError("问答模型未返回文本回答")
        updated_state["answer"] = answer
        return updated_state


if __name__ == "__main__":
    import asyncio
    import json

    from pydantic import BaseModel

    from app.processor.nodes.input_node import InputNode

    async def main():
        # 直接测试问答，仅需 LLM 配置
        state: QueryState = {
            "request": {
                "prompt": "销售额统计的展示可以怎么优化？先给建议，不要修改",
                "screen": {
                    "canvas": {"width": 1920, "height": 1080},
                    "nodes": [
                        {
                            "id": "sales-chart",
                            "type": "bar-chart",
                            "name": "销售额统计",
                            "layout": {"x": 26, "y": 34.8, "width": 400, "height": 260},
                            "props": {"option": {"title": {"text": "销售额统计"}}},
                        }
                    ],
                },
                "context": {"selected_node_ids": ["sales-chart"]},
            }
        }
        state = InputNode().process(state)
        node = QANode()
        final_state = await node.process(state)
        print(
            json.dumps(
                {
                    key: value.model_dump(mode="json")
                    if isinstance(value, BaseModel)
                    else value
                    for key, value in final_state.items()
                },
                ensure_ascii=False,
                indent=2,
            )
        )

    asyncio.run(main())