"""需要物料字段的模型调用：模型按需调用字段查询工具，最后按指定结构输出

动作生成和修正建议共用；字段查询走物料字段 MCP 服务，服务不可用时直接报错
"""

from typing import TypeVar

from langchain.agents import create_agent
from langchain.agents.middleware import wrap_model_call
from langchain.agents.structured_output import ToolStrategy
from langchain_core.messages import AIMessage, HumanMessage
from pydantic import BaseModel

from app.clients.llm import get_llm
from app.clients.material_mcp import get_material_tools
from app.config import settings
from app.processor.state import QueryState
from app.resources.component import get_catalog_cache
from app.resources.material_fields import list_field_groups

OutputModel = TypeVar("OutputModel", bound=BaseModel)


@wrap_model_call
async def _limit_field_queries(request, handler):
    """查询轮数用尽后只保留输出工具，模型必须给出结果"""
    rounds = sum(isinstance(message, AIMessage) for message in request.messages)
    if rounds >= settings.max_field_tool_rounds:
        request = request.override(tools=[])
    return await handler(request)


async def invoke_with_field_tools(
    schema: type[OutputModel], system_prompt: str, content: str
) -> OutputModel | None:
    """返回结构化结果；模型未给出结果时返回 None

    结果不符合 schema 时抛出 StructuredOutputValidationError，由调用方转换为校验错误
    """
    agent = create_agent(
        get_llm(),
        await get_material_tools(),
        system_prompt=system_prompt,
        # 不在调用内自动重试，结构错误交给后续校验和修正流程，与不查字段时一致
        response_format=ToolStrategy(schema, handle_errors=False),
        middleware=[_limit_field_queries],
        name=f"{schema.__name__}_with_fields",
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=content)]},
        # 每轮模型调用和工具执行各占一步，轮数由中间件限制，这里只防止异常循环
        config={"recursion_limit": 2 * settings.max_field_tool_rounds + 6},
    )
    return result.get("structured_response")


def field_group_directory(state: QueryState) -> dict[str, dict]:
    """本次相关组件的字段分组目录，按物料类型列出，省去模型先查目录的一轮调用"""
    nodes = {node.id: node.type for node in state["input_context"].screen.nodes}
    plan = state.get("screen_plan")
    if plan is None:
        types = [
            nodes[target]
            for change in state["requirement"].changes
            if change.target_scope == "nodes"
            for target in change.target_ids
            if target in nodes
        ]
    else:
        types = [
            item.component_type
            if item.disposition in ("add", "replace")
            else nodes.get(item.target_id)
            for item in plan.components
            if item.disposition in ("add", "replace", "modify")
        ]
    catalog = get_catalog_cache()
    directory = {}
    for component_type in dict.fromkeys(item for item in types if item):
        try:
            if catalog is None:
                raise ValueError("物料目录未载入")
            directory[component_type] = {
                key: value
                for key, value in list_field_groups(catalog, component_type).items()
                if key not in ("component_type", "catalog_version")
            }
        except ValueError as error:
            # 没有物料定义的组件不能修改属性，告诉模型原因
            directory[component_type] = {"error": str(error)}
    return directory