"""
工具调用退化纠偏器（Tool Call Degradation Corrector）
====================================================

目的：
- 解决模型/网关在特定条件下，把结构化 tool_calls 退化为 `<tool_call>...</tool_call>` 文本的问题。
- 将退化文本解析为 LangChain 可执行的结构化 `tool_calls`，从而让 Agent 继续运行而不中断。

设计约束：
- 仅作为内部能力在 OpenCode-style Agent 中使用，对外不导出公共接口。
- 必须做工具白名单校验：只允许调用已注册/已注入的工具。

说明：
- 该纠偏器以“包装器”的方式工作：对外暴露 invoke/ainvoke，并把其他属性/方法透传给底层 llm。
- DeepAgent/中间件一般以 Runnable 方式使用 model，因此 duck typing 足够。
"""

from __future__ import annotations

import json
import logging
import re
import uuid
from typing import Any, Dict, List, Optional, Sequence, Set

from langchain.agents.structured_output import OutputToolBinding
from langgraph.prebuilt.tool_node import ToolNode
from langchain_core.messages import AIMessage

logger = logging.getLogger(__name__)


_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>(.*?)</tool_call>", re.DOTALL)
# 兼容另一类常见退化格式，例如：
# CALL_TOOL: {"name": "read_file", "file_path": "/foo"}
_CALL_TOOL_JSON_RE = re.compile(
    r"(?:CALL_TOOL|TOOL_CALL)\s*:\s*(\{.*?\})",
    re.DOTALL,
)


class ToolCallDegradationCorrector:
    """
    包装一个 ChatModel，当检测到 `<tool_call>...</tool_call>` 退化输出时，进行纠偏。

    纠偏触发条件（严格）：
    - response.content 包含 `<tool_call>` 标签
    - 且 response.tool_calls 为空 / 不存在
    """

    def __init__(
        self,
        llm: Any,
        allowed_tool_names: Set[str],
    ) -> None:
        self._llm = llm
        self._allowed_tool_names = allowed_tool_names

    # ----------------------------
    # 属性透传（例如 llm.streaming）
    # ----------------------------
    @property
    def streaming(self) -> Any:
        return getattr(self._llm, "streaming", None)

    @streaming.setter
    def streaming(self, value: Any) -> None:
        setattr(self._llm, "streaming", value)

    def __getattr__(self, name: str) -> Any:
        # 兜底透传：DeepAgent/中间件可能访问底层 llm 的其他方法/属性（如 bind_tools 等）
        return getattr(self._llm, name)

    def bind_tools(
        self,
        tools: Sequence[Any],
        *,
        tool_choice: str | None = None,
        **kwargs: Any,
    ) -> "ToolCallDegradationCorrector":
        """
        DeepAgent/LangChain 会在运行时调用 model.bind_tools(...) 生成“绑定工具后的模型”。

        如果这里直接透传到底层 llm，会导致 wrapper 丢失（返回的是底层 llm），从而纠偏逻辑失效。
        因此必须：
        - 先对底层 llm 执行 bind_tools
        - 再把返回的 bound model 重新包装为 ToolCallDegradationCorrector
        """
        bound = self._llm.bind_tools(tools, tool_choice=tool_choice, **kwargs)
        return ToolCallDegradationCorrector(
            llm=bound,
            allowed_tool_names=self._allowed_tool_names,
        )

    # ----------------------------
    # Runnable 接口：invoke / ainvoke
    # ----------------------------
    def invoke(self, input: Any, **kwargs: Any) -> Any:
        resp = self._llm.invoke(input, **kwargs)
        return self._correct_if_needed(resp)

    async def ainvoke(self, input: Any, **kwargs: Any) -> Any:
        resp = await self._llm.ainvoke(input, **kwargs)
        return self._correct_if_needed(resp)

    # ----------------------------
    # 核心纠偏逻辑
    # ----------------------------
    def _correct_if_needed(self, resp: Any) -> Any:
        if not isinstance(resp, AIMessage):
            return resp

        tool_calls: Optional[Any] = getattr(resp, "tool_calls", None)
        if tool_calls:
            # 正常结构化 tool_calls，不需要纠偏
            return resp

        content = resp.content
        if not isinstance(content, str):
            return resp

        # 支持多种退化格式，从文本中尽力解析 tool_calls
        parsed_tool_calls = self._parse_text_tool_calls(content)
        if not parsed_tool_calls:
            return resp

        cleaned_content = _TOOL_CALL_TAG_RE.sub("", content).strip()

        logger.warning(
            "检测到工具调用退化，已纠偏为结构化 tool_calls | tools=%s",
            [tc.get("name") for tc in parsed_tool_calls],
        )

        return self._clone_ai_message_with_tool_calls(
            original=resp,
            content=cleaned_content,
            tool_calls=parsed_tool_calls,
        )

    def _parse_text_tool_calls(self, content: str) -> List[Dict[str, Any]]:
        """
        从 content 中解析一个或多个 <tool_call>...</tool_call>。

        支持两类 JSON：
        - {"name": "...", "args": {...}}
        - {"name": "...", "param1": "...", ...}  # 扁平化参数
        """
        matches: List[str] = []

        # 1) <tool_call>...</tool_call> 包裹的 JSON
        matches.extend(_TOOL_CALL_TAG_RE.findall(content))

        # 2) 行内 CALL_TOOL: {...} 或 TOOL_CALL: {...} 形式
        matches.extend(_CALL_TOOL_JSON_RE.findall(content))

        if not matches:
            return []

        tool_calls: List[Dict[str, Any]] = []
        for raw in matches:
            raw = raw.strip()
            if not raw:
                continue
            try:
                payload = json.loads(raw)
            except Exception as e:  # noqa: BLE001
                logger.warning("解析退化 <tool_call> JSON 失败：%s | raw=%r", e, raw[:300])
                continue

            if not isinstance(payload, dict):
                continue

            tool_name = payload.get("name")
            if not tool_name or not isinstance(tool_name, str):
                continue

            # 白名单校验：只允许调用已注册工具
            if tool_name not in self._allowed_tool_names:
                logger.warning(
                    "退化工具调用包含未注册工具，已忽略 | tool=%s | allowed=%s",
                    tool_name,
                    sorted(self._allowed_tool_names),
                )
                continue

            if "args" in payload and isinstance(payload.get("args"), dict):
                tool_args = payload["args"]
            else:
                tool_args = {k: v for k, v in payload.items() if k != "name"}

            tool_calls.append(
                {
                    "name": tool_name,
                    "args": tool_args,
                    "id": f"degraded-corrected-{uuid.uuid4().hex[:12]}",
                    "type": "tool_call",
                }
            )

        return tool_calls

    def _clone_ai_message_with_tool_calls(
        self,
        *,
        original: AIMessage,
        content: str,
        tool_calls: List[Dict[str, Any]],
    ) -> AIMessage:
        """
        生成一个新的 AIMessage，保留原消息的元信息，替换 content/tool_calls。
        """
        # AIMessage 字段随 langchain_core 版本可能变化，这里尽量“保守拷贝”
        # 只拷贝常见字段，避免传入未知字段导致构造失败。
        kwargs: Dict[str, Any] = {}

        for key in ("additional_kwargs", "response_metadata", "usage_metadata", "id", "name"):
            if hasattr(original, key):
                kwargs[key] = getattr(original, key)

        # 如果发生退化并且我们已经成功恢复了结构化 tool_calls，
        # 需要同步修正 finish_reason，否则上层 OpenCode-style Agent 会因为
        # finish_reason="stop" 误判为结束而不去执行工具。
        #
        # 这里约定：被纠偏的消息等价于 "tool-calls" 场景。
        rm = kwargs.get("response_metadata") or {}
        if not isinstance(rm, dict):
            rm = {}
        rm["finish_reason"] = "tool-calls"
        kwargs["response_metadata"] = rm

        # invalid_tool_calls 也一起保留（如果存在）
        if hasattr(original, "invalid_tool_calls"):
            kwargs["invalid_tool_calls"] = getattr(original, "invalid_tool_calls")

        return AIMessage(
            content=content,
            tool_calls=tool_calls,
            **kwargs,
        )


def _build_allowed_tool_names_for_degradation_correction(
    tool_node: ToolNode | None,
    built_in_tools: list[dict[str, Any] | Any],
    structured_output_tools: dict[str, OutputToolBinding[Any]],
) -> Set[str]:
    """Build tool name whitelist for ToolCallDegradationCorrector."""
    allowed_tool_names: Set[str] = set()

    # 1) Tools executed via ToolNode (client-side BaseTool / callables)
    if tool_node:
        allowed_tool_names.update(tool_node.tools_by_name.keys())

    # 2) Provider-side tools (dict format) passed through to the LLM
    for t in built_in_tools:
        name_attr = t.get("name") if isinstance(t, dict) else None
        if isinstance(name_attr, str):
            allowed_tool_names.add(name_attr)

    # 3) Structured output tools derived from response_format / AutoStrategy
    allowed_tool_names.update(structured_output_tools.keys())

    return allowed_tool_names

