"""工具参数兼容：模型把列表、对象参数写成 JSON 字符串时，按工具的参数 schema 解析回来

部分模型（例如 claude-haiku-5.5）在嵌套较深、内容较长时，常把 apply_actions 的 actions 整个写成 JSON 字符串，
或把列表里的每一项写成字符串；参数校验报 "Input should be a valid list" 后它往往原样再交，评测里出现过
连续 10 次直到调用次数用尽。这类问题和哪个工具无关，所以在工具调用入口统一处理，对本地工具和 MCP 工具都生效

- 合法的输入不动：schema 里任何一处接受字符串（包括联合类型的某个成员、数组某个成员的 items）时不解析；
  布尔 schema、解析不了的 $ref 这类看不准的情况按"什么都接受"处理，交给原有的参数校验
- 只处理两层：参数本身，以及列表参数里的每一项
- 写成字符串的长 JSON 多数只是结尾的括号多写或少写了一个：原文以括号结尾、模型输出没有因为长度被截断时，
  按还没闭合的括号重新补齐结尾；截断时丢掉的内容补不回来，补齐反而会得到一个意思不同的合法动作，所以不补
- 仍然解析不出来时不执行，用中文说明错在哪里：参数校验的英文报错只说"应该是列表"，模型会照旧写成字符串
- 解析或补齐后执行时，在工具结果末尾说明：补齐只能猜结尾，评测里模型把中间的 option 提前闭合、结尾又少一个括号，
  补齐后 series 落到了 props 下，执行结果对不上而它看不出原因，反复原样提交
"""

import json
from dataclasses import replace
from typing import Any

from langchain.agents.middleware import wrap_tool_call
from langchain_core.messages import AIMessage, ToolMessage
from langchain_core.tools import BaseTool
from langgraph.types import Command
from loguru import logger

STRUCTURED = {"array", "object"}
ANY = frozenset({"string", "number", "integer", "boolean", "null", "array", "object"})
CLOSERS = {"[": "]", "{": "}"}
# 各家接口表示"输出达到长度上限被截断"的写法
TRUNCATED = {"length", "max_tokens"}


def _resolve(schema: Any, defs: dict) -> Any:
    """展开 $ref；指向的定义找不到时按 True（什么都接受）处理"""
    while isinstance(schema, dict) and "$ref" in schema:
        name = schema["$ref"].rsplit("/", 1)[-1]
        if name not in defs:
            return True
        schema = defs[name]
    return schema


def _types(schema: Any, defs: dict) -> set[str]:
    """schema 接受的 JSON 类型：anyOf、oneOf 取并集，allOf 取交集，没有类型约束时什么都接受"""
    schema = _resolve(schema, defs)
    if schema is False:
        return set()
    if not isinstance(schema, dict):
        return set(ANY)
    kind = schema.get("type")
    if kind:
        types = set(kind) if isinstance(kind, list) else {kind}
    elif "properties" in schema:
        types = {"object"}
    elif "items" in schema:
        types = {"array"}
    else:
        types = set(ANY)
    for member in schema.get("allOf", []):
        types &= _types(member, defs)
    if members := [*schema.get("anyOf", []), *schema.get("oneOf", [])]:
        types &= set().union(*(_types(member, defs) for member in members))
    return types


def _items(schema: Any, defs: dict) -> list:
    """各个数组成员的 items schema；有数组成员没有限定 items 时，其中有一个 True（什么都接受）"""
    schema = _resolve(schema, defs)
    if not isinstance(schema, dict):
        return [True] if schema is True else []
    kind = schema.get("type")
    own = set(kind) if isinstance(kind, list) else {kind}
    found = []
    if "prefixItems" in schema:
        # 每个位置各有 schema 的元组，不逐项解析
        found.append(True)
    elif "items" in schema:
        found.append(schema["items"])
    elif "array" in own:
        found.append(True)
    for key in ("anyOf", "oneOf", "allOf"):
        for member in schema.get(key, []):
            found.extend(_items(member, defs))
    return found


def _reclose(text: str) -> str | None:
    """去掉结尾连续的 ] 和 }，按还没闭合的括号重新补齐；结尾停在字符串中间时补不了"""
    body = text.rstrip().rstrip("]} \t\r\n")
    stack: list[str] = []
    in_string = escaped = False
    for char in body:
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
        elif char == '"':
            in_string = True
        elif char in CLOSERS:
            stack.append(char)
        elif char in "]}" and (not stack or CLOSERS[stack.pop()] != char):
            # 中间就有括号对不上，不是结尾的问题，补不了
            return None
    if in_string or not stack:
        return None
    return body + "".join(CLOSERS[item] for item in reversed(stack))


def _matches(value: Any, types: set[str]) -> bool:
    return (isinstance(value, list) and "array" in types) or (
        isinstance(value, dict) and "object" in types
    )


def _parse(value: Any, types: set[str], repair: bool) -> tuple[Any, str | None, bool]:
    """schema 只接受列表或对象、却收到字符串时解析成 JSON；返回 (值, 解析不了时的说明, 是否补了结尾括号)"""
    if not isinstance(value, str) or "string" in types or not types & STRUCTURED:
        return value, None, False
    expected = "数组" if "array" in types else "对象"
    patched = False
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as error:
        # 原文不以括号结尾多半是被截断了，补齐结尾会丢掉后面的内容，不补
        ends_closed = value.rstrip().endswith(("]", "}"))
        repaired = _reclose(value) if repair and ends_closed else None
        try:
            parsed, patched = (json.loads(repaired) if repaired else None), True
        except json.JSONDecodeError:
            parsed = None
        if not _matches(parsed, types):
            problem = (
                f"应该直接写成 JSON {expected}，不要写成字符串；"
                f"这个字符串也不是合法的 JSON（第 {error.pos} 个字符附近：{error.msg}）"
            )
            return value, problem, False
    if not _matches(parsed, types):
        return value, f"应该直接写成 JSON {expected}，不要写成字符串", False
    return parsed, None, patched


def normalize_args(
    tool: BaseTool, args: dict, *, repair: bool = True
) -> tuple[dict, list[str], list[str]]:
    """按工具的参数 schema 把写成字符串的列表、对象参数解析回来

    返回新的参数（不修改传入的）、解析不了的参数说明、补了结尾括号的参数名；有说明时不该执行这次调用
    repair=False 时不补括号，用于模型输出被截断的情况
    """
    schema = tool.tool_call_schema
    schema = schema if isinstance(schema, dict) else schema.model_json_schema()
    defs = schema.get("$defs", {})
    properties = schema.get("properties", {})
    fixed, problems, repaired = {}, [], []
    for name, value in args.items():
        prop = properties.get(name)
        if prop is None:
            fixed[name] = value
            continue
        value, problem, patched = _parse(value, _types(prop, defs), repair)
        if problem:
            problems.append(f"{name} {problem}")
        item_schemas = _items(prop, defs)
        if isinstance(value, list) and item_schemas:
            item_types = set().union(*(_types(item, defs) for item in item_schemas))
            parsed_items = []
            for index, item in enumerate(value):
                item, problem, item_patched = _parse(item, item_types, repair)
                if problem:
                    problems.append(f"{name} 的第 {index + 1} 项{problem}")
                patched = patched or item_patched
                parsed_items.append(item)
            value = parsed_items
        if patched:
            repaired.append(name)
        fixed[name] = value
    return fixed, problems, repaired


def _truncated(state: Any, call_id: str) -> bool:
    """发起这次调用的模型输出是不是因为长度上限被截断"""
    messages = state.get("messages", []) if isinstance(state, dict) else []
    message = next(
        (
            item
            for item in reversed(messages)
            if isinstance(item, AIMessage)
            and any(call["id"] == call_id for call in item.tool_calls)
        ),
        None,
    )
    metadata = message.response_metadata if message is not None else {}
    return bool(
        {metadata.get("finish_reason"), metadata.get("stop_reason")} & TRUNCATED
    )


@wrap_tool_call
async def normalize_tool_args(request, handler):
    """在参数校验前统一修正参数；没有改动时原样执行，解析不了时返回中文说明，不执行"""
    tool = request.tool
    args = request.tool_call.get("args")
    if tool is None or not isinstance(args, dict):
        return await handler(request)
    truncated = _truncated(request.state, request.tool_call["id"])
    fixed, problems, repaired = normalize_args(tool, args, repair=not truncated)
    if problems:
        logger.info(
            f"工具 {tool.name} 的参数写成了不合法的 JSON 字符串"
            f"{'（模型输出被截断）' if truncated else ''}: {problems}"
        )
        return ToolMessage(
            "参数格式错误，这次调用没有执行：\n"
            + "\n".join(f"- {item}" for item in problems)
            + ("\n- 输出达到了长度上限被截断，分成几批提交" if truncated else ""),
            tool_call_id=request.tool_call["id"],
            name=tool.name,
            status="error",
        )
    if fixed == args:
        return await handler(request)
    changed = [name for name in args if fixed[name] != args[name]]
    strict = [name for name in changed if name not in repaired]
    notes = []
    if strict:
        logger.info(
            f"工具 {tool.name} 的参数 {'、'.join(strict)} 写成了 JSON 字符串，已解析"
        )
        notes.append(f"{'、'.join(strict)} 写成了 JSON 字符串，已解析后执行")
    if repaired:
        logger.info(
            f"工具 {tool.name} 的参数 {'、'.join(repaired)} 写成了 JSON 字符串，结尾括号对不上，补齐后解析"
        )
        # 补的是结尾，括号多写或少写在中间时，补齐后层级是错的，但它是合法的 JSON，执行时不一定报错
        notes.append(
            f"{'、'.join(repaired)} 写成了 JSON 字符串，结尾的括号数量不对，已按还没闭合的括号补齐后执行；"
            "如果中间的括号写错了，补齐后的层级也是错的，对照执行结果核对"
        )
    result = await handler(
        request.override(tool_call={**request.tool_call, "args": fixed})
    )
    # 不告诉模型时它会一直写成字符串，补齐猜错层级时它也看不出原因
    return _with_note(
        result, "参数说明：" + "；".join(notes) + "；以后直接写成 JSON，不要写成字符串"
    )


def _with_note(result: Any, note: str) -> Any:
    """在工具返回的消息末尾加一句说明；工具返回 Command 时改其中的工具消息"""
    if isinstance(result, ToolMessage):
        return result.model_copy(update={"content": f"{result.content}\n{note}"})
    if isinstance(result, Command) and isinstance(result.update, dict):
        messages = [
            item.model_copy(update={"content": f"{item.content}\n{note}"})
            if isinstance(item, ToolMessage)
            else item
            for item in result.update.get("messages", [])
        ]
        return replace(result, update={**result.update, "messages": messages})
    return result