"""将简单修改需求转换为动作列表"""

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.action import ACTION_SYSTEM_PROMPT
from app.schemas.actions import ActionBatch


def can_generate_actions(state: QueryState) -> bool:
    requirement = state["requirement"]
    scope = state["change_scope"]
    # 当前只接入三种已有组件操作，后续能力逐步扩展
    return (
        requirement.intent == "modify_screen"
        and not requirement.needs_clarification
        and not scope.needs_planning
        and not scope.resource_needs
        and bool(requirement.changes)
        and all(
            change.target_scope == "nodes"
            and change.operation
            in (
                "move",
                "resize",
                "update",
            )
            for change in requirement.changes
        )
    )


class ActionNode(BaseNode):
    name = "action_node"

    async def process(self, state: QueryState) -> QueryState:
        if not can_generate_actions(state):
            raise ValueError(
                "ActionNode 当前只支持无需规划和资源的 move、resize、update 需求"
            )
        batch = await get_llm(ActionBatch).ainvoke(
            [
                SystemMessage(content=ACTION_SYSTEM_PROMPT),
                HumanMessage(
                    content=json.dumps(
                        {
                            "input_context": state["input_context"].model_dump(
                                mode="json"
                            ),
                            "requirement": state["requirement"].model_dump(mode="json"),
                            "change_scope": state["change_scope"].model_dump(
                                mode="json"
                            ),
                        },
                        ensure_ascii=False,
                    )
                ),
            ]
        )
        batch = ActionBatch.model_validate(batch.model_dump())
        allowed = {}
        for change in state["requirement"].changes:
            for target_id in change.target_ids:
                allowed.setdefault(target_id, set()).add(change.operation)
        existing = {node.id for node in state["input_context"].screen.nodes}
        for action in batch.actions:
            if action.target_id not in existing:
                raise ValueError(f"动作引用了不存在的组件: {action.target_id}")
            if action.type not in allowed.get(action.target_id, set()):
                raise ValueError(
                    f"动作超出需求范围: {action.target_id} / {action.type}"
                )
        expected = {
            (target, operation) for target, ops in allowed.items() for operation in ops
        }
        actual = {(action.target_id, action.type) for action in batch.actions}
        if expected - actual:
            raise ValueError(
                f"动作遗漏了需求中的目标或操作: {sorted(expected - actual)}"
            )
        return {
            **state,
            "actions": batch.actions,
            "validation_passed": None,
            "validation_errors": [],
            "simulated_screen": None,
        }


if __name__ == "__main__":
    import asyncio

    from pydantic import BaseModel

    from app.processor.nodes.input_node import InputNode
    from app.schemas.change_scope import ChangeScope
    from app.schemas.requirement import Requirement

    async def main():
        state = InputNode().process(
            {
                "request": {
                    "prompt": "将销售额统计向右移动 50px",
                    "screen": {
                        "canvas": {"width": 1920, "height": 1080},
                        "nodes": [
                            {
                                "id": "sales-chart",
                                "type": "chart",
                                "layout": {
                                    "x": 100,
                                    "y": 200,
                                    "width": 400,
                                    "height": 300,
                                },
                            }
                        ],
                    },
                }
            }
        )
        state["requirement"] = Requirement(
            intent="modify_screen",
            goal="向右移动 50px",
            needs_clarification=False,
            changes=[
                {
                    "target_scope": "nodes",
                    "target_ids": ["sales-chart"],
                    "operation": "move",
                    "parameters": {"delta_x": 50},
                }
            ],
        )
        state["change_scope"] = ChangeScope(
            scope="local",
            target_ids=["sales-chart"],
            needs_planning=False,
            reason="明确的组件位移",
        )
        node = ActionNode()
        final_state = await node.process(state)
        print(
            json.dumps(
                final_state,
                ensure_ascii=False,
                indent=2,
                default=lambda value: (
                    value.model_dump(mode="json")
                    if isinstance(value, BaseModel)
                    else str(value)
                ),
            )
        )

    asyncio.run(main())