"""Middleware: блокирует вызовы несуществующих инструментов."""
from __future__ import annotations

from typing import Any

from langchain.agents.middleware import AgentMiddleware, AgentState
from langchain_core.messages import ToolMessage
from langchain_core.tools import BaseTool


class SanitizeToolCallsMiddleware(AgentMiddleware[AgentState[Any], Any]):
    """Перехватывает вызовы инструментов с неизвестными именами и возвращает
    понятное сообщение об ошибке вместо падения рантайма."""

    def __init__(self, known_tools: set[str]):
        self.known_tools = known_tools
        self.tools_by_name: dict[str, BaseTool] = {}

    @staticmethod
    def _normalize_tool_name(name: str) -> str:
        # gpt-oss иногда шлёт read_file<|channel|>commentary вместо read_file
        if "<|channel|>" in name:
            name = name.split("<|channel|>", 1)[0]
        return name.strip()

    def _prepare_request(self, request) -> str:
        """Нормализует имя и перепривязывает tool после lookup в ToolNode."""
        raw_name = request.tool_call.get("name", "")
        name = self._normalize_tool_name(raw_name)
        if name != raw_name:
            request.tool_call["name"] = name
        if request.tool is None and name in self.tools_by_name:
            request.tool = self.tools_by_name[name]
        return name

    def _reject(self, request, name: str) -> ToolMessage:
        available = sorted(self.known_tools)
        return ToolMessage(
            content=(
                f"Инструмент '{name}' не существует. "
                f"Доступные инструменты: {available}. "
                "Исправь имя инструмента и попробуй снова."
            ),
            tool_call_id=request.tool_call["id"],
            name=name,
        )

    def wrap_tool_call(self, request, handler):
        name = self._prepare_request(request)
        if name not in self.known_tools:
            return self._reject(request, name)
        if request.tool is None:
            return self._reject(request, name)
        return handler(request)

    async def awrap_tool_call(self, request, handler):
        name = self._prepare_request(request)
        if name not in self.known_tools:
            return self._reject(request, name)
        if request.tool is None:
            return self._reject(request, name)
        return await handler(request)