"""Dynamic discovery of MCP tools, resources, and prompts.

This module implements `DynamicToolRegistry`:
- Initialize MCP server via `initialize`/`initialized`
- Discover `tools/list`, `resources/list`, and `prompts/list`
- Handle `notifications/tools/list_changed` by triggering re-discovery
- Merge registries across multiple MCP servers with a configurable strategy
- Convert discovered tools to provider-specific tool schema formats
"""

from __future__ import annotations

import asyncio
import copy
import json
import logging
import uuid
from dataclasses import dataclass
from typing import Any, Literal

from mcpbridge.config import ENV_KEY_MAP
from mcpbridge.exceptions import ConfigValidationError, MCPProtocolError
from mcpbridge.transport.base import BaseTransport

logger = logging.getLogger("mcpbridge")


@dataclass(slots=True)
class MCPTool:
    """Normalized representation of an MCP tool."""

    name: str
    description: str
    input_schema: dict[str, Any]
    server_name: str
    namespaced_name: str


def _deepcopy_json_schema(schema: dict[str, Any]) -> dict[str, Any]:
    return copy.deepcopy(schema)


def _strip_for_gemini(schema: dict[str, Any]) -> dict[str, Any]:
    """Remove JSON Schema keys that Gemini often rejects."""
    schema = _deepcopy_json_schema(schema)

    def _strip(obj: Any) -> Any:
        if isinstance(obj, dict):
            obj.pop("$schema", None)
            obj.pop("$defs", None)
            # Gemini doesn't like additionalProperties in many schemas.
            obj.pop("additionalProperties", None)
            for v in list(obj.values()):
                _strip(v)
        elif isinstance(obj, list):
            for v in obj:
                _strip(v)
        return obj

    return _strip(schema)


def _json_schema_to_cohere_parameter_definitions(schema: dict[str, Any]) -> dict[str, Any]:
    """Convert JSON Schema tool args to Cohere's flat parameter_definitions format."""
    schema = schema or {}
    properties = schema.get("properties") or {}
    required = set(schema.get("required") or [])
    out: dict[str, Any] = {}

    for prop_name, prop_schema in properties.items():
        if not isinstance(prop_schema, dict):
            continue
        p_type = prop_schema.get("type", "string")
        description = prop_schema.get("description", "")
        out[prop_name] = {
            "description": description,
            "type": p_type,
            "required": prop_name in required,
        }

    return out


class DynamicToolRegistry:
    """In-memory registry for MCP tools/resources/prompts."""

    def __init__(
        self,
        *,
        server_name: str,
        namespace_strategy: Literal["prefix", "error", "last_wins"] = "prefix",
    ) -> None:
        self.server_name = server_name
        self.namespace_strategy = namespace_strategy

        self.tools: dict[str, MCPTool] = {}  # keyed by namespaced_name
        self.resources: dict[str, dict[str, Any]] = {}  # keyed by uri
        self.prompts: dict[str, dict[str, Any]] = {}  # keyed by (possibly namespaced) prompt name

        self._conversion_cache: dict[str, list[dict[str, Any]]] = {}

        self._discover_task: asyncio.Task[None] | None = None

    async def discover(self, transport: BaseTransport) -> None:
        """Initialize and discover all tools/resources/prompts for this server."""
        await self._initialize(transport)

        # Tools
        tools_resp = await self._call(transport, "tools/list", {})
        tools_list = _extract_list(tools_resp, result_key="tools")
        if not tools_list:
            logger.info("MCP tools/list returned empty list for %s", self.server_name)
            self.tools = {}
        else:
            self.tools = {}
            for t in tools_list:
                if not isinstance(t, dict):
                    continue
                name = str(t.get("name", ""))
                if not name:
                    continue
                desc = str(t.get("description", ""))
                input_schema = t.get("inputSchema") or t.get("input_schema") or {}
                if not isinstance(input_schema, dict):
                    input_schema = {}
                namespaced_name = self._namespace_tool_name(name)
                tool = MCPTool(
                    name=name,
                    description=desc,
                    input_schema=input_schema,
                    server_name=self.server_name,
                    namespaced_name=namespaced_name,
                )
                self.tools[namespaced_name] = tool

        # Resources
        resources_resp = await self._call(transport, "resources/list", {})
        resources_list = _extract_list(resources_resp, result_key="resources")
        self.resources = {}
        for r in resources_list or []:
            if not isinstance(r, dict):
                continue
            uri = r.get("uri")
            if uri:
                self.resources[str(uri)] = r

        # Prompts
        prompts_resp = await self._call(transport, "prompts/list", {})
        prompts_list = _extract_list(prompts_resp, result_key="prompts")
        self.prompts = {}
        for p in prompts_list or []:
            if not isinstance(p, dict):
                continue
            name = p.get("name")
            if not name:
                continue
            name = str(name)
            # Prompts are also namespaced to avoid collisions across servers.
            prompt_key = self._namespace_prompt_name(name)
            self.prompts[prompt_key] = p

        # Start a notification watcher (one watcher per transport).
        # We rely on transport implementations to expose `notifications_queue`.
        if hasattr(transport, "notifications_queue"):
            self._start_list_changed_watcher(transport)

    async def _initialize(self, transport: BaseTransport) -> None:
        protocol_versions = ["2024-11-05", "2025-11-25"]
        last_err: BaseException | None = None
        for pv in protocol_versions:
            try:
                init_payload = {
                    "jsonrpc": "2.0",
                    "id": str(uuid.uuid4()),
                    "method": "initialize",
                    "params": {
                        "protocolVersion": pv,
                        "capabilities": {
                            "roots": {"tools": True, "resources": True, "prompts": True},
                            "sampling": {},
                            "tools": {"listChanged": True},
                        },
                        "clientInfo": {"name": "mcpbridge", "version": "0.2.0"},
                    },
                }
                resp = await transport.send(init_payload)
                if isinstance(resp, dict) and "error" in resp:
                    err = resp.get("error") or {}
                    code = err.get("code")
                    msg = err.get("message")
                    # If protocol version is unsupported, try next candidate.
                    if code in (-32602, -32001, -32603) or ("protocolVersion" in str(msg)) or ("protocol version" in str(msg)):
                        supported = err.get("data", {}).get("supportedVersions") if isinstance(err.get("data"), dict) else None
                        if supported:
                            # Try intersection with our candidates.
                            remaining = [v for v in protocol_versions if v in supported]
                            if remaining:
                                continue
                        raise MCPProtocolError(f"Unsupported MCP protocolVersion: {pv}: {msg}")

                # `initialized` notification.
                await transport.send_notification("initialized", {})
                return
            except Exception as e:  # noqa: BLE001
                last_err = e
                logger.warning("MCP initialize failed for protocolVersion=%s (%s). Retrying...", pv, e)
        if last_err:
            raise MCPProtocolError(f"Failed to initialize MCP server for all protocol versions: {last_err}") from last_err

    def _namespace_tool_name(self, tool_name: str) -> str:
        if self.namespace_strategy == "prefix":
            return f"{self.server_name}__{tool_name}"
        return tool_name

    def _namespace_prompt_name(self, prompt_name: str) -> str:
        if self.namespace_strategy == "prefix":
            return f"{self.server_name}__{prompt_name}"
        return prompt_name

    async def _call(self, transport: BaseTransport, method: str, params: dict[str, Any]) -> dict[str, Any]:
        payload = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": method, "params": params}
        resp = await transport.send(payload)
        if isinstance(resp, dict) and resp.get("error"):
            raise MCPProtocolError(f"MCP call {method} failed: {resp['error']}")
        return resp

    def _start_list_changed_watcher(self, transport: BaseTransport) -> None:
        # Fire-and-forget. The watcher will run until the process exits or transport is closed.
        if self._discover_task and not self._discover_task.done():
            return

        async def _watch() -> None:
            q = getattr(transport, "notifications_queue", None)
            if q is None:
                return
            while True:
                msg = await q.get()
                try:
                    if not isinstance(msg, dict):
                        continue
                    method = msg.get("method")
                    if method == "notifications/tools/list_changed":
                        logger.info("MCP tools list_changed for %s -> re-discovering", self.server_name)
                        await self.discover(transport)
                except Exception as e:  # noqa: BLE001
                    logger.warning("Failed to handle list_changed (%s)", e)

        self._discover_task = asyncio.create_task(_watch(), name=f"mcpbridge-registry-watcher-{self.server_name}")

    def merge(self, other: "DynamicToolRegistry") -> None:
        """Merge tools/resources/prompts from another registry."""
        # Merge strategy uses this registry's namespace_strategy.
        for tool_key, tool in other.tools.items():
            if tool_key not in self.tools:
                self.tools[tool_key] = tool
                continue

            # Collision handling.
            existing = self.tools[tool_key]
            if existing.server_name == tool.server_name:
                # Identical source; ignore.
                continue

            if self.namespace_strategy == "prefix":
                raise ConfigValidationError(f"Tool collision after namespacing for `{tool_key}`")
            if self.namespace_strategy == "error":
                raise ConfigValidationError(f"Tool collision across MCP servers for `{tool_key}`")
            if self.namespace_strategy == "last_wins":
                self.tools[tool_key] = tool

        for uri, res in other.resources.items():
            self.resources[uri] = res

        for name, p in other.prompts.items():
            # Same namespacing behavior as tools.
            self.prompts[name] = p

        self._conversion_cache = {}

    def get_tool(self, name: str) -> MCPTool:
        """Lookup tool by namespaced name or raw name."""
        if name in self.tools:
            return self.tools[name]
        # Fallback to raw name match.
        for tool in self.tools.values():
            if tool.name == name:
                return tool
        raise KeyError(f"Tool not found: {name}")

    def for_llm(self, provider: str) -> list[dict[str, Any]]:
        """Convert MCP tools into the given provider's tool schema."""
        if provider in self._conversion_cache:
            return self._conversion_cache[provider]

        tools = list(self.tools.values())
        if not tools:
            self._conversion_cache[provider] = []
            return []

        if provider == "anthropic":
            out = [
                {"name": t.namespaced_name, "description": t.description, "input_schema": t.input_schema}
                for t in tools
            ]
        elif provider in {"openai", "groq", "together", "azure_openai", "openai_compatible", "mistral", "ollama"}:
            out = [
                {
                    "type": "function",
                    "function": {
                        "name": t.namespaced_name,
                        "description": t.description,
                        "parameters": t.input_schema,
                    },
                }
                for t in tools
            ]
        elif provider == "gemini":
            decls = [
                {
                    "name": t.namespaced_name,
                    "description": t.description,
                    "parameters": _strip_for_gemini(t.input_schema),
                }
                for t in tools
            ]
            out = [{"function_declarations": decls}]
        elif provider == "cohere":
            out = [
                {
                    "name": t.namespaced_name,
                    "description": t.description,
                    "parameter_definitions": _json_schema_to_cohere_parameter_definitions(t.input_schema),
                }
                for t in tools
            ]
        elif provider == "bedrock":
            # Bedrock is a special-case wrapper.
            out = [
                {"toolSpec": {"name": t.namespaced_name, "description": t.description, "inputSchema": {"json": t.input_schema}}}
                for t in tools
            ]
        else:
            raise ConfigValidationError(f"Unsupported provider for tool conversion: {provider}")

        self._conversion_cache[provider] = out
        return out


def _extract_list(resp: dict[str, Any], *, result_key: str) -> list[Any] | None:
    """Extract a list from an MCP JSON-RPC response."""
    if not isinstance(resp, dict):
        return None
    result = resp.get("result")
    if isinstance(result, dict) and isinstance(result.get(result_key), list):
        return result.get(result_key)
    if isinstance(result, list):
        return result
    return None

