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

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

from collections.abc import Sequence
from dataclasses import dataclass, field
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 pydantic import Field

from app.actions.batch import BatchError, run_batch
from app.agent.data_tools import DataContext
from app.agent.query_values import EmbeddedQuery, encode_problems, fill_queries
from app.schemas.actions import Action, AddAction, AgentAction
from app.schemas.component_resource import ComponentCatalog
from app.schemas.screen import Screen


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


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(
    before: Screen,
    after: Screen,
    actions: list[Action],
    undeclared_sample: frozenset[str] = frozenset(),
    queries: Sequence[EmbeddedQuery] = (),
) -> str:
    """执行结果：除了动作数量，还给出执行后的画布和每个组件的布局

    模型提交后才第一次看到结果，能自己发现位置超出画布、互相重叠、没排整齐等问题再修正，
    不用靠重复提交来试；本次动过布局或新增的组件前面标 *，已删除的单独列出
    """
    lines = [f"执行成功，共 {len(actions)} 个动作已生效"]
    sample = [
        action.node.id
        for action in actions
        if isinstance(action, AddAction) and action.sample_data
    ]
    # 模型没写 sample_data 时是后端按物料猜的；说明这一点和怎么改，否则模型只看到“用了示例数据”，
    # 以为是自己的数据没生效，会反复提交同样的内容
    if guessed := [item for item in sample if item in undeclared_sample]:
        lines.append(
            f"未写 sample_data，按物料默认标成了示例数据：{'、'.join(guessed)}；"
            "如果这些组件展示的数据都来自大屏上已有的组件或上传的数据文件，把 sample_data 设为 false 重新提交整批，否则不用处理"
        )
    if declared := [item for item in sample if item not in undeclared_sample]:
        lines.append(f"标为示例数据的新增组件：{'、'.join(declared)}")
    # 模型提交时只写了 SQL，看不到结果；给出列名和行数，用来核对 encode 和数据量
    if queries:
        lines.append("$query 已替换为查询结果：")
        for item in queries:
            note = "（结果为空，组件会没有数据）" if not item.rows else ""
            lines.append(
                f"- 第 {item.index + 1} 个动作 {item.node_id} 的 {item.path}："
                f"{item.rows} 行{note}，列 {'、'.join(item.columns)}"
            )

    old_layouts = {node.id: node.layout for node in before.nodes}
    lines.append(
        f"执行后画布 {after.canvas.width:g}x{after.canvas.height:g}，组件布局（* 为本次改动）："
    )
    for node in after.nodes:
        layout = node.layout
        mark = "*" if old_layouts.get(node.id) != layout else " "
        lines.append(
            f"{mark} {node.id} {node.type} "
            f"x={layout.x:g} y={layout.y:g} width={layout.width:g} height={layout.height:g}"
        )
    if removed := [
        item for item in old_layouts if item not in {n.id for n in after.nodes}
    ]:
        lines.append(f"已删除：{'、'.join(removed)}")
    return "\n".join(lines)


@tool
def apply_actions(
    actions: Annotated[
        list[AgentAction],
        # 空的一批会清掉之前成功的结果；参数校验失败不改状态，之前的结果保留
        Field(min_length=1, description="至少一个动作；不需要修改时不要调用，直接回复"),
    ],
    runtime: ToolRuntime[ScreenAgentContext],
) -> Command:
    """提交本次要对大屏做的全部修改

    每次调用都是完整的一批，从原始大屏开始按顺序模拟执行，不在上一次调用的结果上累加
    成功时返回执行后的画布和各组件布局，用它核对结果，不对再重新提交整批
    失败时整批不生效，返回带动作序号的错误，修正后重新提交整批；以最后一次调用的结果为准
    """
    try:
        actions, queries = fill_queries(actions, runtime.context.data)
        completed, after = run_batch(
            runtime.context.screen, runtime.context.catalog, actions
        )
        problems = encode_problems(after, {item.node_id for item in queries})
    except BatchError as error:
        problems = [str(error)]
    if problems:
        message = "执行失败，整批未生效：\n" + "\n".join(
            f"- {item}" for item in problems
        )
        return Command(
            update={
                "actions": None,
                "apply_error": message,
                "messages": [
                    ToolMessage(
                        message, tool_call_id=runtime.tool_call_id, status="error"
                    )
                ],
            }
        )
    # 和上一次成功的提交完全一样：结果早已生效，再提交不会有变化，直接告诉模型结束，不再空转
    if completed == runtime.state.get("actions"):
        return Command(
            update={
                "messages": [
                    ToolMessage(
                        "这一批的结果和上一次提交完全相同，已经生效，不要重复提交；"
                        "任务完成就调用 AgentReply 回复用户，没完成就改成不同的内容",
                        tool_call_id=runtime.tool_call_id,
                    )
                ],
            }
        )
    undeclared = frozenset(
        action.node.id
        for action in actions
        if isinstance(action, AddAction)
        and "sample_data" not in action.model_fields_set
    )
    return Command(
        update={
            "actions": completed,
            "apply_error": None,
            "messages": [
                ToolMessage(
                    _result_text(
                        runtime.context.screen, after, completed, undeclared, queries
                    ),
                    tool_call_id=runtime.tool_call_id,
                )
            ],
        }
    )