"""screen agent 的本地工具：读取当前大屏、在当前大屏上提交一批修改

它们读写的是本次请求的大屏和 agent 状态，所以放在本地；物料等能力由 MCP 提供
读取、修改和校验都基于同一份草稿：原始大屏加上之前写入草稿的修改；回复 done 时草稿交付给用户
"""

import json
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 AIMessage, BaseMessage, ToolMessage
from langgraph.types import Command
from pydantic import Field

from app.actions.batch import BatchError, misplaced_props, run_batch
from app.agent.data_tools import DataContext
from app.agent.query_values import (
    EmbeddedQuery,
    encode_problems,
    fill_queries,
    rebound_nodes,
)
from app.resources.material_fields import paragraph_texts, retained_examples
from app.schemas.actions import (
    Action,
    AddAction,
    AgentAction,
    DeleteAction,
    UpdateAction,
)
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)  # 本会话可用的数据表


# 两个字段只由 apply_actions 写，同一轮只执行第一次提交，所以一步里最多一次写入，不需要合并规则
# 提交是否失败不另存字段，由 last_write_failed 从消息推出：参数解析、参数校验、查询、执行各处的失败都只体现在消息上
class ScreenAgentState(AgentState):
    # 草稿：最近一次写入后的大屏；还没有写入时不存在，用原始大屏
    working_screen: NotRequired[Screen]
    # 交付记录：各次写入的动作（补全、展开后）按顺序拼接，在原始大屏上依次执行得到草稿
    actions: NotRequired[list[Action]]


def current_screen(runtime: ToolRuntime[ScreenAgentContext]) -> Screen:
    return runtime.state.get("working_screen") or runtime.context.screen


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

    读取的是草稿：原始大屏加上之前写入草稿的修改；和 apply_actions 在同一轮调用时读不到这次提交的结果，要核对就在下一轮读取
    """
    nodes = {node.id: node for node in current_screen(runtime).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],
    catalog: ComponentCatalog | None = None,
    queries: Sequence[EmbeddedQuery] = (),
    unchanged: bool = False,
    repeated: bool = False,
) -> str:
    """执行结果：除了动作数量，还给出执行后的画布和每个组件的布局

    模型提交后才第一次看到结果，能自己发现位置超出画布、互相重叠、没排整齐等问题再修正；
    带显示条件的组件标出条件，不同时显示的组件重叠不算问题；
    这一批动过布局或新增的组件前面标 *，这一批删除的单独列出
    没有带来变化时也照样给出布局：只说"已经写入"会让模型以为达到了目标，可能当前大屏本来就是这样、并不符合要求
    统一说"写入草稿"、"回复 done 时交付"：说"已生效"时模型回复 unable 也会说已经改了什么，
    不说怎么交付时它会再交一批空的或重复的来"确认"
    """
    if unchanged and repeated:
        # 之前的提交已经改过：模型不确定是否写入才会重复提交，要写明已经在草稿里、不要重复，并说怎么结束，
        # 只说"没有变化"时它以为修改没成功，会换着字段反复提交直到调用次数用尽
        lines = [
            (
                "执行成功，但这一批没有改变草稿（0 个动作）：提交的内容和草稿里已有的一样，之前的提交已经改过。"
                "对照要求和下面的布局：已经达到就调用 AgentReply 回复 done，草稿会一起交付；"
                "没达到就提交不同的修改，不要原样重复提交"
            )
        ]
    elif unchanged:
        # 这次运行还没有改过大屏：原始大屏本来就是这样，模型的方案并没有改动任何东西，只陈述事实；
        # 加上"完成就回复"之类的提示时，大屏并不符合要求它也会当成完成
        lines = [
            "执行成功，共 0 个动作写入草稿：提交的修改和原始大屏已有的一样，大屏原本就是这样"
        ]
    else:
        lines = [
            f"执行成功，共 {len(actions)} 个动作已写入草稿；回复 done 时草稿交付给用户，不需要再提交确认"
        ]
    if declared := [
        action.node.id
        for action in actions
        if isinstance(action, AddAction) and action.sample_data
    ]:
        lines.append(
            f"声明为自己编写数据的新增组件，交付时标为示例数据：{'、'.join(declared)}"
        )
    # 新增组件没写到的字段沿用默认配置，其中的标题、系列名、数据是演示内容；不说出来时模型以为都按自己的意思写好了，
    # 结果真实数据的卡片带着示例的前缀和环比、图例还叫"访问量"。只陈述事实和后果，由模型决定改不改
    if examples := _retained_example_notes(actions, after, catalog):
        lines.append(
            "新增组件里这些值还是物料的示例内容：示例文字和组件要展示的内容不符就改写；"
            "示例数据会让组件交付时标为示例数据，要展示真实数据就换掉（数据来自上传的文件时用 $query），"
            "没有真实数据时可以保留并在回复里说明"
        )
        lines.extend(examples)
    # 问答之后"把结果放到大屏上"，模型会把分析结论抄成文本组件；skill 里写了规则仍会发生，写入时再陈述一次，
    # 用户明确要在大屏上放这段文字时由模型保留。要说清空出的位置怎么办：只说"改成简短的标题"时，
    # 模型把结论块缩成一行小标题，留下一大块空位，反复加删组件去填，直到调用次数用尽
    if paragraphs := _paragraph_notes(actions, after, catalog):
        lines.append(
            "这一批写入的标题或文字是一段话：大屏上的文字只放标题和简短标签，"
            "分析结论、洞察、数据说明和注意事项写在回复里。用户明确要求在大屏上放这段文字时保留；"
            "否则只放这段文字的文本组件直接删除，原来的位置给图表或让相邻组件占满，"
            "图表标题、指标名这类改成简短的名称"
        )
        lines.extend(paragraphs)
    # 模型提交时只写了 SQL，看不到结果；给出列名和行数，用来核对 encode 和数据量
    if queries:
        lines.append("$query 已替换为查询结果：")
        for item in queries:
            note = "（结果为空，组件会没有数据）" if not item.rows else ""
            result = (
                f"取值 {item.value}"
                if item.value is not None
                else f"{item.rows} 行{note}，列 {'、'.join(item.columns)}"
            )
            lines.append(
                f"- 第 {item.index + 1} 个动作 {item.node_id} 的 {item.path}：{result}"
            )

    old_layouts = {node.id: node.layout for node in before.nodes}
    # 整页切换时不同页的组件放在同一块区域，标出显示条件，模型才不会把它们当成重叠去挪开
    lines.append(
        f"执行后画布 {after.canvas.width:g}x{after.canvas.height:g}，组件布局（* 为本次改动）："
    )
    if any(node.visibleWhen is not None for node in after.nodes):
        lines.append(
            "标了「仅 变量=值 时显示」的组件只在变量等于该值时出现，"
            "同一变量取不同值的组件不会同时显示，它们之间重叠不算问题"
        )
    for node in after.nodes:
        layout = node.layout
        mark = "*" if old_layouts.get(node.id) != layout else " "
        condition = (
            f" 仅 {node.visibleWhen.var}={json.dumps(node.visibleWhen.equals, ensure_ascii=False)} 时显示"
            if node.visibleWhen is not None
            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}"
            + condition
        )
    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)


def _retained_example_notes(
    actions: Sequence[Action], after: Screen, catalog: ComponentCatalog | None
) -> list[str]:
    """这一批新增的组件上还和默认配置一样的示例内容，每个组件一行"""
    nodes = {node.id: node for node in after.nodes}
    notes = []
    for action in actions:
        if not isinstance(action, AddAction) or action.node.id not in nodes:
            continue
        node = nodes[action.node.id].model_dump(mode="json")
        if retained := retained_examples(catalog, node):
            notes.append(
                f"- {action.node.id}："
                + "、".join(
                    f"{field.path}（{field.label}）{_brief(value)}"
                    for field, value in retained
                )
            )
    return notes


def _paragraph_notes(
    actions: Sequence[Action], after: Screen, catalog: ComponentCatalog | None
) -> list[str]:
    """这一批写入的识别字段里成段的文字，每个组件一行

    新增的组件看全部识别字段，修改只看这一批写到的字段：只改颜色时不提原有的说明文字
    """
    written: dict[str, set[str] | None] = {}
    for action in actions:
        if isinstance(action, AddAction):
            written[action.node.id] = None
        elif (
            isinstance(action, UpdateAction)
            and written.get(action.target_id, set()) is not None
        ):
            written.setdefault(action.target_id, set()).add(action.field_path)
    nodes = {node.id: node for node in after.nodes}
    notes = []
    for node_id, paths in written.items():
        if node_id not in nodes:
            continue
        node = nodes[node_id].model_dump(mode="json")
        if found := paragraph_texts(catalog, node, paths):
            notes.append(
                f"- {node_id}："
                + "、".join(
                    f"{field.path}（{field.label}）{len(text.strip())} 字 {_brief(' '.join(text.split()))}"
                    for field, text in found
                )
            )
    return notes


def _brief(value) -> str:
    if isinstance(value, list):
        return f"示例数据 {len(value)} 行"
    if isinstance(value, str):
        return f"“{value if len(value) <= 20 else value[:20] + '…'}”"
    return json.dumps(value, ensure_ascii=False)


def _reply(
    runtime: ToolRuntime, text: str, *, error: bool = False, **update
) -> Command:
    """返回工具消息；状态字段由调用方显式传入，没传的保持不变"""
    status = "error" if error else "success"
    message = ToolMessage(text, tool_call_id=runtime.tool_call_id, status=status)
    return Command(update={**update, "messages": [message]})


def _first_write(message: AIMessage) -> str | None:
    """一条模型消息里第一个 apply_actions 调用的 ID；同一轮只执行它"""
    return next(
        (call["id"] for call in message.tool_calls if call["name"] == "apply_actions"),
        None,
    )


def _is_first_submission(runtime: ToolRuntime) -> bool:
    """按模型消息里 tool_calls 的顺序，这次调用是不是第一个 apply_actions

    同一轮的多个调用都从同一份草稿出发，都执行会互相覆盖；只执行第一个，与并发谁先跑完无关
    直接调用工具、找不到发起调用的模型消息时，视为第一个
    """
    message = next(
        (
            item
            for item in reversed(runtime.state.get("messages", []))
            if isinstance(item, AIMessage)
            and any(call["id"] == runtime.tool_call_id for call in item.tool_calls)
        ),
        None,
    )
    return message is None or _first_write(message) == runtime.tool_call_id


def last_write_failed(messages: Sequence[BaseMessage]) -> bool:
    """最后一次被接受的提交是否失败：最后一条带 apply_actions 的模型消息里，第一个调用的结果是错误

    参数解析、参数校验、查询、执行任何一处失败，结果都是 status="error" 的工具消息，都在这里判断，
    不靠各处各自记状态，以后多一种失败也不会漏；同一轮被拒绝的后续调用不是第一个，不算
    没有结果时也算失败：这一批没有执行，最后想做的修改没有写入草稿
    """
    for position in range(len(messages) - 1, -1, -1):
        message = messages[position]
        if isinstance(message, AIMessage) and (call_id := _first_write(message)):
            reply = next(
                (
                    item
                    for item in messages[position + 1 :]
                    if isinstance(item, ToolMessage) and item.tool_call_id == call_id
                ),
                None,
            )
            return reply is None or reply.status == "error"
    return False


@tool
def apply_actions(
    actions: Annotated[
        list[AgentAction],
        # 不限制至少一个：以前整批替换时空的一批会清掉之前的结果，增量执行下空的一批只是没有变化，
        # 按普通的无变化回复，比参数校验的英文报错更能让模型知道该结束了
        Field(description="这一批要执行的动作；不需要修改时不要调用，直接回复"),
    ],
    runtime: ToolRuntime[ScreenAgentContext],
) -> Command:
    """把一批修改写进这次的草稿；回复 done 时草稿交付给用户，不需要再提交确认

    第一次提交从原始大屏开始，之后每一批都在草稿上继续修改；回复 unable 或 clarify 时草稿整个丢弃
    成功时返回这一批执行后的画布和各组件布局，用它核对结果；结果不对再提交一批修正，已经写入的动作不用重复提交
    一批里任何一个动作失败，这一批都不写入草稿，之前写入的保留；按带序号的错误修正这一批后再提交
    同一轮只执行第一次调用，同一轮的 get_nodes 也读不到这次的结果；看到执行结果后再调用 AgentReply
    """
    if not _is_first_submission(runtime):
        # 被拒绝的调用不写任何状态字段，不会覆盖第一次提交的结果
        return _reply(
            runtime,
            "同一轮只执行第一次 apply_actions，这一批没有执行；需要的话在下一轮根据执行结果再提交",
            error=True,
        )
    before = current_screen(runtime)
    misplaced = misplaced_props(actions, runtime.context.catalog)
    # 写错层级的属性不生效，失败和成功时都说明；它常是 encode 对不上的真正原因
    hint = (
        "\n新增组件写了物料没有的属性，前端不会使用："
        + "；".join(misplaced)
        + "；检查是不是写错了层级"
        if misplaced
        else ""
    )
    try:
        actions, queries = fill_queries(
            actions, runtime.context.data, runtime.context.catalog, before
        )
        completed, after = run_batch(before, runtime.context.catalog, actions)
        # 这一批改过数据或映射的组件都核对，不论数据来自这一批的 $query、之前的提交还是直接写入
        problems = encode_problems(after, rebound_nodes(before, after))
    except BatchError as error:
        problems = [str(error)]
    if problems:
        # 失败只体现在工具消息上，由 last_write_failed 判断，不写状态字段
        message = "执行失败，这一批未写入草稿，之前写入的修改保留：\n" + "\n".join(
            f"- {item}" for item in problems
        )
        return _reply(runtime, message + hint, error=True)
    log = runtime.state.get("actions", [])
    delivered = [*log, *completed]
    # 按交付结果判断有没有变化，不比较参数：同样的 move 再执行一次会继续移动，同样的 update 再执行一次没有变化
    # 交付结果除了大屏还有只在动作上的"自己编写的数据"声明，只比大屏会把"重新 add 改 sample_data"当成没变化而丢掉
    unchanged = after == before and _sample_marks(delivered) == _sample_marks(log)
    # 这次运行之前已有成功提交时，没有变化多半是在重复自己的修改；还没有时说明原始大屏本来就是这样
    text = _result_text(
        before,
        after,
        completed,
        runtime.context.catalog,
        queries,
        unchanged,
        repeated=bool(log),
    )
    if unchanged:
        return _reply(runtime, text + hint)
    return _reply(runtime, text + hint, working_screen=after, actions=delivered)


def _sample_marks(actions: Sequence[Action]) -> dict[str, bool]:
    """按交付记录，每个新增组件是否声明为自己编写的数据；这个声明只在 add 动作上，不在大屏里"""
    marks: dict[str, bool] = {}
    for action in actions:
        if isinstance(action, AddAction):
            marks[action.node.id] = action.sample_data
        elif isinstance(action, DeleteAction):
            marks.pop(action.target_id, None)
    return marks