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

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

from typing import Any, Literal, Self

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

from app.agent.data_tools import DATA_TOOLS
from app.agent.skills import skills_prompt
from app.agent.tools import (
    ScreenAgentContext,
    ScreenAgentState,
    apply_actions,
    get_nodes,
)
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.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：无法完成"
    )
    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


# 模型没有调用工具时最多纠正几次
NO_TOOL_CALL_RETRIES = 2
NO_TOOL_CALL_HINT = (
    "每一轮都要调用工具。要结束时调用 AgentReply，给用户的话写在 message 里"
)


@wrap_model_call
async def require_tool_call(request, handler):
    """ToolStrategy 要求每一轮都调用工具，结束时调用 AgentReply；模型偶尔只回文本（想直接回复用户）或空回复

    原样重试往往得到同样的结果，或者模型转而重复提交 apply_actions 直到调用次数用尽，
    所以把它的回复和纠正说明加在下一次调用里；只用于这一次调用，不写进 agent 的消息
    """
    response = await handler(request)
    for _ in range(NO_TOOL_CALL_RETRIES):
        message = response.result[0] if response.result else None
        if not isinstance(message, AIMessage) or message.tool_calls:
            break
        logger.warning(
            f"模型没有调用工具，纠正后重新调用: {str(message.content)[:200]}"
        )
        # 空回复不放回去，部分接口不接受内容为空的助手消息
        replied = [message] if message.content else []
        request = request.override(
            messages=[
                *request.messages,
                *replied,
                HumanMessage(content=NO_TOOL_CALL_HINT),
            ]
        )
        response = await handler(request)
    return response


# 是否挂数据工具 → 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,
            middleware=[
                require_tool_call,
                # 超过上限时插入说明并结束，没有结构化结果，由节点按失败处理
                ModelCallLimitMiddleware(
                    run_limit=settings.max_agent_steps, exit_behavior="end"
                ),
            ],
            # 不继承外层图的检查点：agent 内部没有中断，中间消息不需要持久化
            checkpointer=False,
            name="screen_agent",
        )
    return _agents[with_data]