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

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

from typing import 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
from loguru import logger
from pydantic import BaseModel, Field, model_validator

from app.agent.skills import skill_middleware
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


# 偶发的空回复最多重新调用几次
EMPTY_REPLY_RETRIES = 2


@wrap_model_call
async def retry_empty_reply(request, handler):
    """每一轮都必须调用工具（最后一轮是 AgentReply）；模型偶尔返回既没有文本也没有工具调用的空回复，
    这是调用失败而不是模型的决定，重新调用即可"""
    response = await handler(request)
    for _ in range(EMPTY_REPLY_RETRIES):
        message = response.result[0] if response.result else None
        if not isinstance(message, AIMessage) or message.tool_calls:
            break
        logger.warning(
            f"模型返回了没有工具调用的回复，重新调用: {message.response_metadata}"
        )
        response = await handler(request)
    return response


_agent = None


async def get_screen_agent():
    """第一次使用时组装，之后复用；系统提示是常量，不随物料目录变化"""
    global _agent
    if _agent is None:
        _agent = create_agent(
            model=get_llm(),
            tools=[get_nodes, apply_actions, *await get_mcp_tools()],
            system_prompt=SYSTEM_PROMPT,
            response_format=ToolStrategy(AgentReply),
            state_schema=ScreenAgentState,
            context_schema=ScreenAgentContext,
            middleware=[
                *skill_middleware(),
                retry_empty_reply,
                # 超过上限时插入说明并结束，没有结构化结果，由节点按失败处理
                ModelCallLimitMiddleware(
                    run_limit=settings.max_agent_steps, exit_behavior="end"
                ),
            ],
            # 不继承外层图的检查点：agent 内部没有中断，中间消息不需要持久化
            checkpointer=False,
            name="screen_agent",
        )
    return _agent