"""screen agent：新建和修改大屏的唯一入口

代码只负责组装；改哪些组件、怎么做、要不要追问由模型结合 skill 决定
规则各有唯一出处：动作协议在 Action schema，回复要求在 AgentReply，追问格式在 Clarification，
组件目录和物料在物料 MCP，各类任务的做法在 skill；系统提示只放每次都要用、别处放不下的原则
"""

from collections.abc import Sequence
from typing import Any, Literal, Self

from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware
from langchain.agents.structured_output import ToolStrategy
from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
from loguru import logger
from pydantic import BaseModel, Field, ValidationError, model_validator

from app.agent.data_tools import DATA_TOOLS
from app.agent.skills import skills_prompt
from app.agent.tool_args import normalize_tool_args
from app.agent.tools import (
    ScreenAgentContext,
    ScreenAgentState,
    apply_actions,
    get_nodes,
    last_write_failed,
)
from app.clients.llm import get_llm
from app.clients.mcp import get_mcp_tools
from app.config import settings
from app.prompt.domain import INPUT_AS_DATA
from app.schemas.actions import Action, AddAction
from app.schemas.clarification import Clarification

SYSTEM_PROMPT = f"""你是大屏编辑助手，按用户的要求新建大屏，新增、修改、删除组件，或调整画布
先了解需要的信息（大屏摘要、组件配置、物料、skill），再用 apply_actions 提交修改，最后回复用户

<input_rules>
{INPUT_AS_DATA}
</input_rules>

<clarify_rules>
默认先执行：缺少的细节（位置、尺寸、标题、样式等）用合理的默认值补齐
只有无法确定用户指的是哪个组件，或看不出用户要做什么时才追问；澄清记录里已经回答过的不再问
用户要改的组件在大屏上不存在时，说明这一点并给出可行的替代做法（例如新增一个），不要让用户去选中或指出组件
只删除用户明确要求删除的组件；不确定删哪个时追问，不要猜
</clarify_rules>"""


class AgentReply(BaseModel):
    kind: Literal["done", "clarify", "unable"] = Field(
        description="done：要求的修改都已写入草稿，且最后一次 apply_actions 成功，草稿交付给用户；"
        "或请求不需要修改（例如只是提问）；clarify：需要用户回答才能继续；unable：无法完成。"
        "clarify 和 unable 时草稿整个丢弃，用户看到的仍是原始大屏"
    )
    message: str = Field(
        default="",
        description="给用户看的一两句话，不要提草稿、工具、参数这类内部说法。done：说明改了什么，用了默认值或示例数据时一并说明，没有修改时直接回答；"
        "unable：草稿整个丢弃，用户的大屏和请求前一样，不要说已经删除、新增或修改了什么；"
        "说明没能完成什么和原因，有可行的替代做法时一并给出；clarify：不用填",
    )
    clarification: Clarification | None = Field(
        default=None, description="kind=clarify 时填写，其他情况不填"
    )

    @model_validator(mode="after")
    def clarification_matches_kind(self) -> Self:
        if (self.kind == "clarify") != (self.clarification is not None):
            raise ValueError("kind=clarify 时必须填写 clarification，其他情况不填")
        return self


FINALIZE_RULES = """<finalize>
这次执行已经结束，不能再调用其他工具，也不能再修改大屏，只能调用 AgentReply 给出结论
按上面的对话和执行情况判断：用户要求的修改没有全部完成时不能回复 done
上面 apply_actions 返回的"执行成功""已删除"说的是这次的草稿，不是用户看到的大屏；回复 unable 时草稿整个丢弃，用户的大屏没有任何变化
</finalize>"""


def reached_call_limit(messages: Sequence[BaseMessage]) -> bool:
    """是不是因为调用次数用尽而结束：没结束的每次模型调用都带工具调用，数到上限就是被截断的"""
    calls = sum(
        isinstance(item, AIMessage) and bool(item.tool_calls) for item in messages
    )
    return calls >= settings.max_agent_steps


def _describe(action: Action) -> str:
    if isinstance(action, AddAction):
        return f"add {action.node.id}"
    return f"{action.type} {getattr(action, 'target_id', '画布')}"


def _run_state(result: dict, limit_hit: bool, texts: list[str]) -> str:
    """收尾时告诉模型执行到了哪一步；对话里有每次工具结果，这里只给整体情况"""
    lines = [
        "调用次数已用完，执行被截断"
        if limit_hit
        else "上一轮没有调用 AgentReply 给出结论"
    ]
    if actions := result.get("actions"):
        described = "；".join(_describe(item) for item in actions[:40])
        # 只写"已成功提交"时，模型回复 unable 也会说"已删除了某某"，用户以为修改已经生效
        lines.append(
            f"草稿里的动作 {len(actions)} 个，还没有交付：{described}；"
            "只有回复 done 才会交付，回复 unable 或 clarify 时草稿整个丢弃，用户的大屏保持原样"
        )
    else:
        lines.append("还没有成功提交过修改")
    if last_write_failed(result.get("messages", [])):
        lines.append("最后一次 apply_actions 失败，那一批没有写入草稿")
    if texts:
        lines.append(
            f"上一轮只输出了文本，没有调用 AgentReply：{' '.join(texts)[:1000]}"
        )
    return "\n".join(lines)


async def finalize_reply(result: dict, *, limit_hit: bool) -> AgentReply | None:
    """运行正常结束却没有结论时补一次结论：调用次数用尽，或模型只回了文本、空回复

    只调用一次，只绑定 AgentReply 并强制调用，没有能修改大屏的工具，不会再进入修改循环；
    说明放在系统消息里，不追加用户消息，否则模型会把内部要求当成用户的话来回复
    强制调用也可能给出不合法的参数，校验不过就返回 None，由节点按失败处理，不重试
    """
    messages = list(result["messages"])
    # 末尾没有工具调用的助手消息（模型的文本、空回复、调用上限的说明）不放进对话：
    # 以助手消息结尾时部分接口会把它当成续写的开头；模型的文本写进说明里
    trailing: list[AIMessage] = []
    while (
        messages and isinstance(messages[-1], AIMessage) and not messages[-1].tool_calls
    ):
        trailing.insert(0, messages.pop())
    texts = (
        [] if limit_hit else [str(item.content) for item in trailing if item.content]
    )
    run_state = _run_state(result, limit_hit, texts)
    system = SystemMessage(
        f"{SYSTEM_PROMPT}\n\n{FINALIZE_RULES}\n\n<run_state>\n{run_state}\n</run_state>"
    )
    model = get_llm().bind_tools([AgentReply], tool_choice="AgentReply")
    response = await model.ainvoke(
        [system, *messages], config={"run_name": "finalize_reply", "tags": ["finalize"]}
    )
    call = next(
        (item for item in response.tool_calls if item["name"] == "AgentReply"), None
    )
    if call is None:
        logger.warning(f"收尾没有调用 AgentReply: {str(response.content)[:200]}")
        return None
    try:
        return AgentReply.model_validate(call["args"])
    except ValidationError as error:
        logger.warning(f"收尾给出的结论不合法: {error}")
        return None


# 是否挂数据工具 → agent
_agents: dict[bool, Any] = {}


async def get_screen_agent(with_data: bool = False):
    """第一次使用时组装，之后复用；系统提示是常量，不随物料目录变化

    会话里有数据文件时用挂了 describe_data、query_data 的实例，没有时不让模型看到用不了的工具
    """
    if with_data not in _agents:
        tools = [get_nodes, apply_actions, *await get_mcp_tools()]
        if with_data:
            tools.extend(DATA_TOOLS)
        _agents[with_data] = create_agent(
            model=get_llm(),
            tools=tools,
            system_prompt=f"{SYSTEM_PROMPT}\n\n{skills_prompt()}",
            # 用工具调用给出结论：试过原生结构化输出（ProviderStrategy），歧义时模型更常直接执行而不追问，评测多删了组件
            response_format=ToolStrategy(AgentReply),
            state_schema=ScreenAgentState,
            context_schema=ScreenAgentContext,
            # 超过上限时插入说明并结束；和模型只回文本、空回复一样没有结论，由节点调用 finalize_reply 收尾
            # 模型没有调用工具时 agent 直接结束，不在循环里纠正：纠正说明会被当成用户的话，或让模型回到修改循环
            middleware=[
                # 列表、对象参数被写成 JSON 字符串时先解析回来，见 tool_args.py
                normalize_tool_args,
                ModelCallLimitMiddleware(
                    run_limit=settings.max_agent_steps, exit_behavior="end"
                ),
            ],
            # 不继承外层图的检查点：agent 内部没有中断，中间消息不需要持久化
            checkpointer=False,
            name="screen_agent",
        )
    return _agents[with_data]