"""
Semantic Tool Filter Hook

Pre-call hook that filters MCP tools semantically before LLM inference.
Reduces context window size and improves tool selection accuracy.
"""

from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Optional

from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict

from litellm._logging import verbose_proxy_logger
from litellm.constants import (
    DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL,
    DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD,
    DEFAULT_MCP_SEMANTIC_FILTER_TOP_K,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
    SemanticToolFilterContextWindowError,
)

if TYPE_CHECKING:
    from litellm.caching.caching import DualCache
    from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
        SemanticMCPToolFilter,
    )
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.router import Router


class SemanticToolFilterConfig(TypedDict, total=False):
    enabled: ReadOnly[bool]
    embedding_model: ReadOnly[str]
    top_k: ReadOnly[int]
    similarity_threshold: ReadOnly[float]


def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> str:
    """Cap a CSV of tool names to max_length, dropping any name that does not fit whole."""
    if len(tool_names_csv) <= max_length:
        return tool_names_csv

    head: Final = tool_names_csv[: max_length + 1]
    if "," not in head:
        return ""

    return head.rsplit(",", 1)[0]


class SemanticToolFilterHook(CustomLogger):
    """
    Pre-call hook that filters MCP tools semantically.

    This hook:
    1. Extracts the user query from messages
    2. Filters tools based on semantic similarity to the query
    3. Returns only the top-k most relevant tools to the LLM
    """

    def __init__(self, semantic_filter: "SemanticMCPToolFilter"):
        """
        Initialize the hook.

        Args:
            semantic_filter: SemanticMCPToolFilter instance
        """
        super().__init__()
        self.filter = semantic_filter

        verbose_proxy_logger.debug(
            "Initialized SemanticToolFilterHook with filter: enabled=%s, top_k=%s",
            semantic_filter.enabled,
            semantic_filter.top_k,
        )

    def _should_expand_mcp_tools(self, tools: Iterable[Mapping[str, object]]) -> bool:
        """
        Check if tools contain MCP references with server_url="litellm_proxy".

        Only expands MCP tools pointing to litellm proxy, not external MCP servers.
        """
        from litellm.responses.mcp.litellm_proxy_mcp_handler import (
            LiteLLM_Proxy_MCP_Handler,
        )

        return LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools)

    async def _expand_mcp_tools(
        self,
        tools: Iterable[Mapping[str, object]],
        user_api_key_dict: "UserAPIKeyAuth",
    ) -> list[dict[str, object]]:
        """
        Expand MCP references to actual tool definitions.

        Reuses LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format
        which internally does: parse -> fetch -> filter -> deduplicate -> transform
        """
        from litellm.responses.mcp.litellm_proxy_mcp_handler import (
            LiteLLM_Proxy_MCP_Handler,
        )

        # Parse to separate MCP tools from other tools
        mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)

        if not mcp_tools:
            return []

        # Use single combined method instead of 3 separate calls
        # This already handles: fetch -> filter by allowed_tools -> deduplicate -> transform
        (
            openai_tools,
            _,
        ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format(
            user_api_key_auth=user_api_key_dict, mcp_tools_with_litellm_proxy=mcp_tools
        )

        # Convert Pydantic models to dicts for compatibility
        openai_tools_as_dicts: Final[list[dict[str, object]]] = []
        for tool in openai_tools:
            if hasattr(tool, "model_dump"):
                tool_dict = tool.model_dump(exclude_none=True)
                verbose_proxy_logger.debug(
                    "Converted Pydantic tool to dict: %s -> dict with keys: %s",
                    type(tool).__name__,
                    list(tool_dict.keys()),
                )
                openai_tools_as_dicts.append(tool_dict)
            elif hasattr(tool, "dict"):
                tool_dict = tool.dict(exclude_none=True)
                verbose_proxy_logger.debug("Converted Pydantic tool (v1) to dict: %s -> dict", type(tool).__name__)
                openai_tools_as_dicts.append(tool_dict)
            elif isinstance(tool, dict):
                verbose_proxy_logger.debug("Tool is already a dict with keys: %s", list(tool.keys()))
                openai_tools_as_dicts.append(tool)
            else:
                verbose_proxy_logger.warning("Tool is unknown type: %s, passing as-is", type(tool))
                openai_tools_as_dicts.append(tool)

        verbose_proxy_logger.debug(
            "Expanded %s MCP reference(s) to %s tools (all as dicts)", len(mcp_tools), len(openai_tools_as_dicts)
        )

        return openai_tools_as_dicts

    async def _filter_expanded_tools(
        self,
        data: dict,
        expanded_tools: list[dict[str, object]],
    ) -> list[dict[str, object]]:
        """
        Apply the semantic filter to expanded MCP tool definitions.

        Expanded tools are flat OpenAI function dicts with a top-level
        "name" (see transform_mcp_tool_to_openai_responses_api_tool), so
        filter_tools can name-match them against the semantic router.
        """
        raw_messages: Final = data.get("messages") or data.get("input") or []
        messages: Final = [{"role": "user", "content": raw_messages}] if isinstance(raw_messages, str) else raw_messages
        user_query: Final = self.filter.extract_user_query(messages)
        if not user_query:
            verbose_proxy_logger.debug("No user query found, skipping semantic filter on expanded MCP tools")
            return expanded_tools

        return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools)

    def _selected_tool_names(self, filtered_tools: Sequence[object]) -> list[str]:
        """Names of the semantically selected tools, as produced by the MCP expansion."""
        names: Final = (self.filter._extract_tool_info(tool)[0] for tool in filtered_tools)
        return [name for name in names if name]

    @staticmethod
    def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
        """
        Restrict each litellm_proxy MCP reference to the semantically selected tools.

        The reference block is preserved rather than replaced with expanded tools, so the
        MCP gateway still performs the expansion. That keeps the per-endpoint tool shape
        and tool auto-execution intact. Expansion already applied any caller-supplied
        allowed_tools, so this selection can only narrow a block further.

        Whether an undecidable selection exposes every tool or none is owned by
        SemanticMCPToolFilter.filter_tools, which returns the full set when nothing
        matches; the same policy therefore governs references and plain tools. Passing an
        empty selection through is safe rather than a hidden allow-all: the gateway reads
        the union of every reference's allowed_tools and treats an empty union as unset.
        """
        from litellm.responses.mcp.litellm_proxy_mcp_handler import (
            LiteLLM_Proxy_MCP_Handler,
        )

        return [
            (
                {**tool, "allowed_tools": selected_tool_names}
                if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
                else tool
            )
            for tool in tools
        ]

    def _is_mcp_tool(self, tool: object) -> bool:
        """
        Check whether *tool* is registered in the MCP semantic router.

        Classification strategy (shape-first, lookup-second):
        1. Chat Completions format dicts are always native.
        2. Responses API function tools are always native.
        3. Everything else is looked up by name in the MCP registry.
        """
        if isinstance(tool, dict) and tool.get("type") == "function" and isinstance(tool.get("function"), dict):
            return False
        if isinstance(tool, dict) and tool.get("type") == "function" and isinstance(tool.get("name"), str):
            return False
        name, _ = self.filter._extract_tool_info(tool)
        return bool(name) and name in self.filter._tool_map

    def _get_metadata_variable_name(self, data: dict) -> str:
        if "litellm_metadata" in data:
            return "litellm_metadata"
        return "metadata"

    def _emit_filter_metadata(
        self,
        data: dict,
        mcp_tools: Sequence[object],
        filtered_mcp_tools: Sequence[object],
        native_tools: Sequence[object],
        filtered_tools: Sequence[object],
    ) -> None:
        """
        Emit response-header metadata when MCP tools were filtered.

        Stats report MCP-only counts so downstream consumers see accurate
        semantic filter metrics.  Skips metadata entirely for purely-native
        requests to avoid spurious headers.
        """
        if mcp_tools:
            filter_stats: Final = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}"
            tool_names_csv: Final = self._get_tool_names_csv(filtered_mcp_tools)

            _metadata_variable_name: Final = self._get_metadata_variable_name(data)
            metadata: Final = data.setdefault(_metadata_variable_name, {})
            metadata["litellm_semantic_filter_stats"] = filter_stats
            metadata["litellm_semantic_filter_tools"] = tool_names_csv

            verbose_proxy_logger.info(
                "Semantic tool filter: %s MCP tools (%s native preserved, %s total)",
                filter_stats,
                len(native_tools),
                len(filtered_tools),
            )
        else:
            verbose_proxy_logger.info(
                "Semantic tool filter: all %s tools are native, no MCP filtering applied", len(native_tools)
            )

    def _emit_filter_metadata_safe(
        self,
        data: dict,
        mcp_tools: Sequence[object],
        filtered_mcp_tools: Sequence[object],
        native_tools: Sequence[object],
        filtered_tools: Sequence[object],
    ) -> None:
        """
        Emit filter metadata without letting an emission failure abort the
        already-filtered request.
        """
        try:
            self._emit_filter_metadata(
                data=data,
                mcp_tools=mcp_tools,
                filtered_mcp_tools=filtered_mcp_tools,
                native_tools=native_tools,
                filtered_tools=filtered_tools,
            )
        except Exception as e:
            verbose_proxy_logger.warning(
                "Failed to emit semantic filter metadata: %s",
                e,
                exc_info=True,
            )

    async def async_pre_call_hook(
        self,
        user_api_key_dict: "UserAPIKeyAuth",
        cache: "DualCache",
        data: dict,
        call_type: str,
    ) -> Exception | str | dict | None:
        """
        Filter tools before LLM call based on user query.

        This hook is called before the LLM request is made. It filters the
        tools list to only include semantically relevant tools.
        """
        if call_type not in ("completion", "acompletion", "aresponses"):
            verbose_proxy_logger.debug("Skipping semantic filter for call_type=%s", call_type)
            return None

        tools: Final = data.get("tools")
        if not tools:
            verbose_proxy_logger.debug("No tools in request, skipping semantic filter")
            return None

        if self._should_expand_mcp_tools(tools):
            verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering")

            if not self.filter.enabled:
                verbose_proxy_logger.debug("Semantic filter disabled, leaving MCP references untouched")
                return None

            try:
                native_tools_before_expand = [t for t in tools if not (isinstance(t, dict) and t.get("type") == "mcp")]

                expanded_tools: Final = await self._expand_mcp_tools(tools, user_api_key_dict)

                if not expanded_tools:
                    verbose_proxy_logger.warning("No tools expanded from MCP references")
                    return None

                filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)

                selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools)
                narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names)
                data["tools"] = narrowed_tools
                self._emit_filter_metadata_safe(
                    data=data,
                    mcp_tools=expanded_tools,
                    filtered_mcp_tools=filtered_expanded_tools,
                    native_tools=native_tools_before_expand,
                    filtered_tools=narrowed_tools,
                )
                verbose_proxy_logger.info(
                    "Expanded MCP references to %s tools (%s native preserved), semantic filter selected %s",
                    len(expanded_tools),
                    len(native_tools_before_expand),
                    len(filtered_expanded_tools),
                )
                return data

            except SemanticToolFilterContextWindowError as e:
                raise HTTPException(status_code=400, detail={"error": str(e)}) from e
            except Exception as e:
                verbose_proxy_logger.error("Failed to expand MCP references: %s", e, exc_info=True)
                return None

        messages = data.get("messages", [])
        if not messages:
            messages = data.get("input", [])
        if not messages:
            verbose_proxy_logger.debug("No messages in request, skipping semantic filter")
            return None

        if not self.filter.enabled:
            verbose_proxy_logger.debug("Semantic filter disabled, skipping")
            return None

        try:
            user_query: Final = self.filter.extract_user_query(messages)
            if not user_query:
                verbose_proxy_logger.debug("No user query found, skipping semantic filter")
                return None

            native_tools: Final[list[object]] = []
            mcp_tools: Final[list[object]] = []
            mcp_indices: Final[set[int]] = set()
            for i, t in enumerate(tools):
                if self._is_mcp_tool(t):
                    mcp_tools.append(t)
                    mcp_indices.add(i)
                else:
                    native_tools.append(t)

            verbose_proxy_logger.debug(
                "Applying semantic filter: %s MCP tools, %s native tools, query: '%s...'",
                len(mcp_tools),
                len(native_tools),
                user_query[:50],
            )

            if mcp_tools:
                filtered_mcp_tools: list[object] = await self.filter.filter_tools(
                    query=user_query,
                    available_tools=mcp_tools,
                )
            else:
                filtered_mcp_tools = []

            filtered_mcp_names: Final[set[str]] = set()
            for t in filtered_mcp_tools:
                name, _ = self.filter._extract_tool_info(t)
                if name:
                    filtered_mcp_names.add(name)

            filtered_tools: Final[list[object]] = []
            for i, t in enumerate(tools):
                if i in mcp_indices:
                    name, _ = self.filter._extract_tool_info(t)
                    if name in filtered_mcp_names:
                        filtered_tools.append(t)
                else:
                    filtered_tools.append(t)

            data["tools"] = filtered_tools

            self._emit_filter_metadata_safe(
                data=data,
                mcp_tools=mcp_tools,
                filtered_mcp_tools=filtered_mcp_tools,
                native_tools=native_tools,
                filtered_tools=filtered_tools,
            )

            return data

        except SemanticToolFilterContextWindowError as e:
            raise HTTPException(status_code=400, detail={"error": str(e)}) from e
        except Exception as e:
            verbose_proxy_logger.warning("Semantic tool filter hook failed: %s. Proceeding with all tools.", e)
            return None

    async def async_post_call_response_headers_hook(
        self,
        data: dict,
        user_api_key_dict: "UserAPIKeyAuth",
        response: object,
        request_headers: dict[str, str] | None = None,
        litellm_call_info: dict[str, object] | None = None,
    ) -> dict[str, str] | None:
        """Add semantic filter stats and tool names to response headers."""
        from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH

        _metadata_variable_name: Final = self._get_metadata_variable_name(data)
        metadata: Final = data.get(_metadata_variable_name, {})

        filter_stats: Final = metadata.get("litellm_semantic_filter_stats")
        if not filter_stats:
            return None

        headers: Final = {"x-litellm-semantic-filter": filter_stats}

        # Add CSV of filtered tool names (nginx-safe length)
        tool_names_csv: Final = metadata.get("litellm_semantic_filter_tools", "")
        header_safe_csv: Final = _truncate_csv_at_tool_name_boundary(
            tool_names_csv=tool_names_csv,
            max_length=MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH,
        )
        if header_safe_csv:
            headers["x-litellm-semantic-filter-tools"] = header_safe_csv

        return headers

    def _get_tool_names_csv(self, tools: Sequence[object]) -> str:
        """Extract tool names and return as CSV string."""
        if not tools:
            return ""

        tool_names: Final = []
        for tool in tools:
            name = tool.get("name", "") if isinstance(tool, dict) else getattr(tool, "name", "")
            if name:
                tool_names.append(name)

        return ",".join(tool_names)

    @staticmethod
    async def initialize_from_config(
        config: SemanticToolFilterConfig | None,
        llm_router: Optional["Router"],
    ) -> Optional["SemanticToolFilterHook"]:
        """
        Initialize semantic tool filter from proxy config.

        Args:
            config: Proxy configuration dict (litellm_settings.mcp_semantic_tool_filter)
            llm_router: LiteLLM router instance for embeddings

        Returns:
            SemanticToolFilterHook instance if enabled, None otherwise
        """
        from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
            SemanticMCPToolFilter,
        )

        if not config or not config.get("enabled", False):
            verbose_proxy_logger.debug("Semantic tool filter not enabled in config")
            return None

        if llm_router is None:
            verbose_proxy_logger.warning("Cannot initialize semantic filter: llm_router is None")
            return None

        try:
            embedding_model: Final = config.get("embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL)
            top_k: Final = config.get("top_k", DEFAULT_MCP_SEMANTIC_FILTER_TOP_K)
            similarity_threshold = config.get("similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD)

            semantic_filter: Final = SemanticMCPToolFilter(
                embedding_model=embedding_model,
                litellm_router_instance=llm_router,
                top_k=top_k,
                similarity_threshold=similarity_threshold,
                enabled=True,
            )

            # Build router from MCP registry on startup
            await semantic_filter.build_router_from_mcp_registry()

            hook: Final = SemanticToolFilterHook(semantic_filter)

            verbose_proxy_logger.info(
                "✅ MCP Semantic Tool Filter enabled: embedding_model=%s, top_k=%s, similarity_threshold=%s",
                embedding_model,
                top_k,
                similarity_threshold,
            )

            return hook

        except ImportError as e:
            verbose_proxy_logger.warning(
                "semantic-router not installed. Install with: pip install 'litellm[semantic-router]'. Error: %s", e
            )
            return None
        except Exception as e:
            verbose_proxy_logger.exception("Failed to initialize MCP semantic tool filter: %s", e)
            return None
