"""screen agent 的本地工具：读取本次请求的大屏、提交整批动作

它们读写的是本次请求的大屏和 agent 状态，所以放在本地；物料等能力由 MCP 提供
"""

from dataclasses import dataclass
from typing import Annotated, NotRequired

from langchain.agents import AgentState
from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from langgraph.types import Command

from app.actions.batch import BatchError, run_batch
from app.schemas.actions import Action, AddAction
from app.schemas.component_resource import ComponentCatalog
from app.schemas.screen import Screen


@dataclass(frozen=True)
class ScreenAgentContext:
    screen: Screen  # 本次请求的原始大屏，只读
    catalog: ComponentCatalog  # 当前物料目录，只用于补全和校验


def _latest(_, new):
    # 同一步里有多次提交时取最后一个，避免并行工具调用写同一字段时报错
    return new


class ScreenAgentState(AgentState):
    # 最近一次 apply_actions 成功时补全后的动作；失败时清空，以最后一次提交为准
    actions: NotRequired[Annotated[list[Action] | None, _latest]]
    # 最近一次 apply_actions 失败的原因，成功时清空
    apply_error: NotRequired[Annotated[str | None, _latest]]


@tool
def get_nodes(node_ids: list[str], runtime: ToolRuntime[ScreenAgentContext]) -> dict:
    """读取组件的完整配置（layout、props、style 等）。修改属性或迁移数据前先读取

    读取的是原始大屏，不包含本次已提交的修改
    """
    nodes = {node.id: node for node in runtime.context.screen.nodes}
    return {
        "nodes": [
            nodes[item].model_dump(mode="json") for item in node_ids if item in nodes
        ],
        "missing": [item for item in node_ids if item not in nodes],
    }


def _result_text(actions: list[Action]) -> str:
    lines = [f"执行成功，共 {len(actions)} 个动作已生效"]
    if sample := [
        action.node.id
        for action in actions
        if isinstance(action, AddAction) and action.sample_data
    ]:
        lines.append(f"使用物料示例数据的新增组件：{'、'.join(sample)}")
    return "\n".join(lines)


@tool
def apply_actions(
    actions: list[Action], runtime: ToolRuntime[ScreenAgentContext]
) -> Command:
    """提交本次要对大屏做的全部修改

    每次调用都是完整的一批，从原始大屏开始按顺序模拟执行，不在上一次调用的结果上累加
    失败时整批不生效，返回带动作序号的错误，修正后重新提交整批；以最后一次调用的结果为准
    """
    try:
        completed, _ = run_batch(
            runtime.context.screen, runtime.context.catalog, actions
        )
    except BatchError as error:
        message = f"执行失败，整批未生效：\n- {error}"
        return Command(
            update={
                "actions": None,
                "apply_error": message,
                "messages": [
                    ToolMessage(
                        message, tool_call_id=runtime.tool_call_id, status="error"
                    )
                ],
            }
        )
    return Command(
        update={
            "actions": completed,
            "apply_error": None,
            "messages": [
                ToolMessage(_result_text(completed), tool_call_id=runtime.tool_call_id)
            ],
        }
    )