"""沿用意图判断，提取具体需求"""

import json

from langchain_core.messages import HumanMessage, SystemMessage

from app.clients.llm import get_llm
from app.processor.base import BaseNode
from app.processor.state import QueryState
from app.prompt.requirement import REQUIREMENT_SYSTEM_PROMPT
from app.schemas.requirement import Requirement


class RequirementNode(BaseNode):
    name = "requirement_node"

    async def process(self, state: QueryState) -> QueryState:
        context = state["input_context"]
        intent = state["intent"]
        payload = {
            "intent": intent,
            "input_context": context.model_dump(mode="json", exclude_unset=True),
        }
        llm = get_llm(Requirement)
        result = await llm.ainvoke(
            [
                SystemMessage(content=REQUIREMENT_SYSTEM_PROMPT),
                HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
            ]
        )

        print(f"LLM 返回的需求: {result}")
        print("=" * 66)

        # 沿用门控后的意图，合并后重新校验字段一致性
        requirement = Requirement.model_validate(
            {**result.model_dump(), "intent": intent}
        )
        node_ids = {node.id for node in context.screen.nodes}
        for change in requirement.changes:
            if unknown_ids := set(change.target_ids) - node_ids:
                raise ValueError(f"需求引用了不存在的节点: {sorted(unknown_ids)}")

        return {**state, "requirement": requirement}


if __name__ == "__main__":
    import asyncio

    from pydantic import BaseModel

    from app.clients.decision import close_decision_engine
    from app.processor.nodes.input_node import InputNode
    from app.processor.nodes.intent_node import IntentNode

    async def main():
        # 1. 修改这里的请求，测试具体需求提取
        state: QueryState = {
            "request": {
                "prompt": "把销售额统计向右移动 50px，其他不变",
                "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": "销售额统计"}}},
                        }
                    ],
                    "dataSource": [],
                },
                "context": {"selected_node_ids": ["sales-chart"]},
            }
        }
        # 2. 依次执行输入、意图判断和需求理解
        try:
            state = InputNode().process(state)
            state = await IntentNode().process(state)
            node = RequirementNode()
            final_state = await node.process(state)
        finally:
            await close_decision_engine()

        # 3. 打印完整 state
        print(
            json.dumps(
                {
                    key: value.model_dump(mode="json", exclude_unset=True)
                    if isinstance(value, BaseModel)
                    else value
                    for key, value in final_state.items()
                },
                ensure_ascii=False,
                indent=2,
            )
        )

    asyncio.run(main())