"""需求理解节点"""

from app.clients.decision import DecisionEngine, get_decision_engine
from app.processor.base import BaseNode
from app.processor.state import QueryState
from app.prompt.requirement import build_intent_question
from app.schemas.base import JsonObject
from app.schemas.decision import DecisionRequest
from app.schemas.input_context import InputContext


class RequirementNode(BaseNode):
    name = "requirement_node"

    def __init__(self, *, decision_engine: DecisionEngine | None = None):
        super().__init__()
        self._decision_engine = decision_engine

    @staticmethod
    def build_decision_state(context: InputContext) -> JsonObject:
        """只整理判断意图需要的摘要，完整 Screen 仍保留在输入上下文"""
        nodes: list[JsonObject] = []
        for node in context.screen.nodes:
            option = node.props.get("option")
            title = option.get("title") if isinstance(option, dict) else None
            titles = title if isinstance(title, list) else [title]
            nodes.append(
                {
                    "id": node.id,
                    "type": node.type,
                    "name": node.name,
                    "titles": [
                        item["text"]
                        for item in titles
                        if isinstance(item, dict) and isinstance(item.get("text"), str)
                    ],
                    "layout": node.layout.model_dump(mode="json", exclude_unset=True),
                }
            )
        return {
            "prompt": context.prompt,
            "context": context.context.model_dump(mode="json", exclude_unset=True),
            "screen_summary": {
                "canvas": {
                    "width": context.screen.canvas.width,
                    "height": context.screen.canvas.height,
                },
                "nodes": nodes,
                "data_sources": [
                    {
                        key: source[key]
                        for key in ("id", "name", "type")
                        if key in source
                    }
                    for source in context.screen.dataSource
                ],
            },
            "attachments": [
                {
                    key: attachment[key]
                    for key in ("id", "kind", "name")
                    if key in attachment
                }
                for attachment in context.attachments
            ],
        }

    async def process(self, state: QueryState) -> QueryState:
        # 构建jev请求
        request = DecisionRequest(
            # 获取 InputNode 整理好的完整上下文
            state=self.build_decision_state(state["input_context"]),
            decisions={"intent": build_intent_question()},
        )

        # 调用 jev
        engine = self._decision_engine or get_decision_engine()
        decision = await engine.run(request)

        updated_state: QueryState = {**state, "intent_decision": decision}
        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():
        # 1. 修改这里的请求，测试不同意图
        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": "销售额统计"}}},
                        }
                    ],
                    "dataSource": [],
                },
                "context": {"selected_node_ids": ["sales-chart"]},
            }
        }
        # 2. 先校验输入，再真实调用 Jev，退出时关闭连接
        state = InputNode().process(state)
        async with DecisionEngine() as engine:
            node = RequirementNode(decision_engine=engine)
            final_state = await node.process(state)
        # 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())