"""意图识别节点"""

from typing import cast

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


class IntentNode(BaseNode):
    name = "intent_node"

    @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,
            "clarification_history": context.clarification_history,
            "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:
        context = state["input_context"]
        # 构建jev请求
        request = DecisionRequest(
            # 将 InputNode 的上下文整理为意图判断摘要
            state=self.build_decision_state(context),
            decisions={"intent": build_intent_question()},
        )

        # 调用 jev
        decision_engine = get_decision_engine()
        # 获取 jev 结果
        decision_res = await decision_engine.run(request)

        # 原生 unclear 或低可信时澄清，原始结果完整保留
        decision = decision_res.results["intent"]
        if (
            decision.choice == "unclear"
            or decision.confidence < self.config.intent_min_confidence
        ):
            intent: Intent = "unclear"
        else:
            intent = cast(Intent, decision.choice)

        updated_state: QueryState = {
            **state,
            "intent_decision": decision_res,
            "intent": intent,
        }
        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，退出时关闭 Jev 连接
        state = InputNode().process(state)
        node = IntentNode()
        try:
            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())