"""MCP Tool Adapter for bridging MCP servers into the ToolRegistry.

Manages the lifecycle of an MCP server connection — handshake, tool
discovery, invocation with schema validation, and clean shutdown.

Operations:
    | Method                        | Description                                |
    |-------------------------------|--------------------------------------------|
    | ``connect()``                 | Start server process / HTTP session        |
    | ``disconnect()``              | Clean shutdown of the server connection    |
    | ``discover_tools(filter)``    | Enumerate tools from the server            |
    | ``invoke(name, arguments)``   | Call a tool with input validation          |
    | ``register_tools(registry)``  | Bulk-register discovered tools             |
"""

from __future__ import annotations

import logging
import threading
import time
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from cleveragents.tool.registry import ToolRegistry

import jsonschema
from pydantic import BaseModel, ConfigDict, Field

from cleveragents.domain.models.core.tool import (
    BindingMode,
    ResourceAccessMode,
    ResourceSlot,
    ToolCapability,
)
from cleveragents.tool.runtime import ToolSpec

logger = logging.getLogger(__name__)

__all__ = [
    "MCPCapabilityMetadata",
    "MCPServerConfig",
    "MCPToolAdapter",
    "MCPToolDescriptor",
    "MCPToolFilter",
    "MCPToolResult",
    "MCPTransport",
]


class MCPServerConfig(BaseModel):
    """Configuration for connecting to an MCP server.

    Attributes:
        name: Server identifier.
        transport: Protocol type (``stdio``, ``sse``, ``streamable-http``).
        command: Command to spawn the server (required for stdio).
        args: Command-line arguments for the server process.
        env: Environment variables passed to the server process.
        url: Server URL (required for sse/streamable-http).
        headers: HTTP headers for remote connections.
    """

    name: str = Field(..., min_length=1, description="Server identifier")
    transport: str = Field(..., description="Transport: stdio | sse | streamable-http")
    command: str | None = Field(default=None, description="Spawn command (stdio)")
    args: list[str] = Field(default_factory=list, description="Command arguments")
    env: dict[str, str] = Field(default_factory=dict, description="Environment vars")
    url: str | None = Field(default=None, description="Server URL (sse/http)")
    headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers")

    model_config = ConfigDict(str_strip_whitespace=True)


class MCPCapabilityMetadata(BaseModel):
    """Full MCP server capability metadata.

    Exposes the complete capability information returned by the MCP
    server during the initialize handshake, including supported
    resources, prompts, and tool schemas.

    Attributes:
        tools: Whether the server supports tool invocation.
        resources: Whether the server supports resource access.
        prompts: Whether the server supports prompt templates.
        raw_capabilities: Full raw capabilities dict from the server.
    """

    tools: bool = Field(default=False, description="Supports tool invocation")
    resources: bool = Field(default=False, description="Supports resource access")
    prompts: bool = Field(default=False, description="Supports prompt templates")
    raw_capabilities: dict[str, Any] = Field(
        default_factory=dict, description="Raw capabilities from server"
    )

    model_config = ConfigDict(str_strip_whitespace=True)


class MCPToolDescriptor(BaseModel):
    """Descriptor for a tool discovered from an MCP server.

    Attributes:
        name: Tool name as exposed by the MCP server.
        description: Human-readable description of the tool.
        input_schema: JSON Schema for tool inputs.
        annotations: Optional tool annotations from the MCP server.
    """

    name: str = Field(..., description="Tool name")
    description: str = Field(default="", description="Tool description")
    input_schema: dict[str, Any] = Field(
        default_factory=dict, description="JSON Schema for inputs"
    )
    annotations: dict[str, Any] = Field(
        default_factory=dict, description="MCP tool annotations"
    )

    model_config = ConfigDict(str_strip_whitespace=True)


class MCPToolResult(BaseModel):
    """Result of an MCP tool invocation.

    Attributes:
        success: Whether the invocation succeeded.
        data: Result payload from the MCP server, normalised to a dict.
            On success the MCP 1.4.0 ``content`` list is wrapped under the
            ``"content"`` key so callers always receive a ``dict[str, Any]``
            regardless of the raw server response format.
        error: Error message on failure.
        duration_ms: Wall-clock execution time in milliseconds.
    """

    success: bool = Field(..., description="Whether invocation succeeded")
    data: dict[str, Any] = Field(default_factory=dict, description="Result payload")
    error: str | None = Field(default=None, description="Error message")
    duration_ms: float = Field(default=0.0, description="Execution time (ms)")

    model_config = ConfigDict(str_strip_whitespace=True)


class MCPToolFilter(BaseModel):
    """Filter controlling which MCP tools are exposed.

    Attributes:
        include: Whitelist of tool names (empty = include all).
        exclude: Blacklist of tool names.
    """

    include: list[str] = Field(default_factory=list, description="Include list")
    exclude: list[str] = Field(default_factory=list, description="Exclude list")


class MCPTransport:
    """Abstract transport layer for MCP server communication.

    Subclass to provide real stdio or HTTP transports. The default
    implementation raises ``NotImplementedError`` for all operations.
    """

    def connect(self, config: MCPServerConfig) -> dict[str, Any]:
        """Perform handshake and return server capabilities."""
        raise NotImplementedError

    def call(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        """Send a JSON-RPC method call and return the result."""
        _ = method, params
        raise NotImplementedError

    def close(self) -> None:
        """Shut down the transport connection."""


class MCPToolAdapter:
    """Adapter bridging an MCP server into the CleverAgents ToolRegistry.

    Parameters
    ----------
    config:
        Server connection configuration.
    transport:
        Optional transport override (for testing). When ``None``, a
        real transport is selected based on ``config.transport``.
    """

    def __init__(
        self,
        config: MCPServerConfig,
        transport: MCPTransport | None = None,
    ) -> None:
        self._validate_config(config)
        self._config = config
        self._transport = transport or MCPTransport()
        self._connected = False
        self._capabilities: dict[str, Any] = {}
        self._tools: dict[str, MCPToolDescriptor] = {}
        self._lock = threading.RLock()
        self._notification_listeners: list[Callable[[str, dict[str, Any]], None]] = []

    @staticmethod
    def _validate_config(config: MCPServerConfig) -> None:
        """Validate transport-specific required fields.

        Raises ``ValueError`` when required fields are missing.
        """
        if config.transport == "stdio" and not config.command:
            msg = (
                f"MCPServerConfig '{config.name}': stdio transport "
                f"requires 'command' field"
            )
            raise ValueError(msg)
        if config.transport in ("sse", "streamable-http") and not config.url:
            msg = (
                f"MCPServerConfig '{config.name}': {config.transport} "
                f"transport requires 'url' field"
            )
            raise ValueError(msg)

    @classmethod
    def from_server_config(
        cls,
        config: MCPServerConfig,
        transport: MCPTransport | None = None,
    ) -> MCPToolAdapter:
        """Create an adapter (validation is performed in ``__init__``)."""
        return cls(config=config, transport=transport)

    @property
    def server_name(self) -> str:
        return self._config.name

    @property
    def transport_type(self) -> str:
        return self._config.transport

    @property
    def is_connected(self) -> bool:
        with self._lock:
            return self._connected

    @property
    def capabilities(self) -> dict[str, Any]:
        with self._lock:
            return dict(self._capabilities)

    @property
    def discovered_tools(self) -> list[MCPToolDescriptor]:
        with self._lock:
            return list(self._tools.values())

    @property
    def capability_metadata(self) -> MCPCapabilityMetadata:
        """Structured capability metadata from the MCP server.

        Returns full capability information parsed from the server's
        initialize handshake response.  Handles both flat capability
        dicts (``{"tools": True}``) and nested ones
        (``{"capabilities": {"tools": True}}``).
        """
        with self._lock:
            raw = dict(self._capabilities)
        # The transport may return capabilities at the top level or nested
        # under a ``capabilities`` key depending on the MCP server.
        caps = raw.get("capabilities", raw)
        return MCPCapabilityMetadata(
            tools=bool(caps.get("tools")),
            resources=bool(caps.get("resources")),
            prompts=bool(caps.get("prompts")),
            raw_capabilities=raw,
        )

    def connect(self, timeout: float = 30.0) -> None:
        """Connect to the MCP server and perform the initialize handshake.

        Parameters
        ----------
        timeout:
            Maximum seconds to wait for the transport handshake.
            Raises ``ConnectionError`` if exceeded. Defaults to 30 s.

        Raises
        ------
        ConnectionError
            If the transport fails to connect or the handshake times out.
        """
        with self._lock:
            if self._connected:
                return

            result: dict[str, Any] = {}
            exc_holder: list[BaseException] = []

            def _do_connect() -> None:
                try:
                    result["caps"] = self._transport.connect(self._config)
                except Exception as exc:
                    exc_holder.append(exc)

            thread = threading.Thread(target=_do_connect, daemon=True)
            thread.start()
            thread.join(timeout=timeout)

            if thread.is_alive():
                msg = (
                    f"MCP connection to '{self._config.name}' timed out "
                    f"after {timeout}s"
                )
                raise ConnectionError(msg)

            if exc_holder:
                msg = f"MCP connection to '{self._config.name}' failed: {exc_holder[0]}"
                raise ConnectionError(msg) from exc_holder[0]

            self._capabilities = result.get("caps") or {}
            self._connected = True
            logger.info("Connected to MCP server '%s'", self._config.name)

    def disconnect(self) -> None:
        """Disconnect from the MCP server."""
        with self._lock:
            if not self._connected:
                return
            try:
                self._transport.close()
            except Exception:
                logger.warning(
                    "Error during MCP disconnect from '%s'",
                    self._config.name,
                    exc_info=True,
                )
            self._connected = False
            self._tools.clear()
            logger.info("Disconnected from MCP server '%s'", self._config.name)

    def reconnect(self, timeout: float = 30.0) -> None:
        """Disconnect from the current session and establish a fresh connection.

        Clears cached tool descriptors. Callers should call
        ``discover_tools()`` again after reconnecting.

        Parameters
        ----------
        timeout:
            Maximum seconds to wait for the new transport handshake.

        Raises
        ------
        ConnectionError
            If the new connection attempt fails or times out.
        """
        self.disconnect()
        self.connect(timeout=timeout)
        logger.info("Reconnected to MCP server '%s'", self._config.name)

    # -- Notification support -------------------------------------------------

    def add_notification_listener(
        self,
        callback: Callable[[str, dict[str, Any]], None],
    ) -> None:
        """Register a callback to receive MCP server notifications.

        The callback is invoked with ``(method, params)`` whenever
        ``dispatch_notification`` is called (e.g. from a transport layer
        or test harness that simulates server-side push events).

        Parameters
        ----------
        callback:
            Callable accepting ``(method: str, params: dict[str, Any])``.
        """
        with self._lock:
            self._notification_listeners.append(callback)

    def dispatch_notification(
        self,
        method: str,
        params: dict[str, Any] | None = None,
    ) -> None:
        """Dispatch an MCP notification to all registered listeners.

        This is the entry point for server-initiated push events such as
        ``notifications/tools/list_changed``.  Transports that support
        push notifications should call this method when a notification
        arrives.  Tests may call it directly to simulate server events.

        Parameters
        ----------
        method:
            The JSON-RPC notification method name.
        params:
            Notification parameters (defaults to empty dict).
        """
        payload = params or {}
        with self._lock:
            listeners = list(self._notification_listeners)

        for listener in listeners:
            try:
                listener(method, payload)
            except Exception:
                logger.warning(
                    "Notification listener raised an exception for method '%s'",
                    method,
                    exc_info=True,
                )

    def discover_tools(
        self,
        tool_filter: MCPToolFilter | None = None,
    ) -> list[MCPToolDescriptor]:
        """Enumerate tools from the connected MCP server.

        Parameters
        ----------
        tool_filter:
            Optional include/exclude filter.

        Raises
        ------
        RuntimeError
            If the adapter is not connected.
        """
        with self._lock:
            if not self._connected:
                msg = (
                    f"MCP adapter '{self._config.name}' is not connected. "
                    f"Call connect() first."
                )
                raise RuntimeError(msg)

            result = self._transport.call("tools/list", {})
            raw_tools = result.get("tools", [])

            descriptors: list[MCPToolDescriptor] = []
            for raw in raw_tools:
                desc = MCPToolDescriptor(
                    name=raw.get("name", ""),
                    description=raw.get("description", ""),
                    input_schema=raw.get("inputSchema", {}),
                    annotations=raw.get("annotations", {}),
                )
                descriptors.append(desc)

            if tool_filter:
                descriptors = self._apply_filter(descriptors, tool_filter)

            self._tools = {d.name: d for d in descriptors}
            return descriptors

    def invoke(
        self,
        tool_name: str,
        arguments: dict[str, Any],
    ) -> MCPToolResult:
        """Invoke a discovered MCP tool by name.

        Validates inputs against the tool's JSON Schema before calling
        the server.

        Parameters
        ----------
        tool_name:
            Name of the tool as discovered from the server.
        arguments:
            Tool input arguments.

        Returns
        -------
        MCPToolResult
            Invocation outcome with success flag, data, and timing.
        """
        with self._lock:
            if not self._connected:
                msg = (
                    f"MCP adapter '{self._config.name}' is not connected. "
                    f"Call connect() first."
                )
                raise RuntimeError(msg)

            descriptor = self._tools.get(tool_name)
            if descriptor is None:
                return MCPToolResult(
                    success=False,
                    error=f"Tool '{tool_name}' not found in MCP server "
                    f"'{self._config.name}'",
                )

            validation_error = self._validate_input(descriptor, arguments)
            if validation_error:
                return MCPToolResult(
                    success=False,
                    error=f"Input validation failed for '{tool_name}': "
                    f"{validation_error}",
                )

            start = time.monotonic()
            try:
                result = self._transport.call(
                    "tools/call",
                    {"name": tool_name, "arguments": arguments},
                )
            except TimeoutError as exc:
                elapsed = (time.monotonic() - start) * 1000
                return MCPToolResult(
                    success=False,
                    error=f"Tool '{tool_name}' timeout: {exc}",
                    duration_ms=elapsed,
                )
            except Exception as exc:
                elapsed = (time.monotonic() - start) * 1000
                return MCPToolResult(
                    success=False,
                    error=f"MCP server error invoking '{tool_name}': {exc}",
                    duration_ms=elapsed,
                )
            elapsed = (time.monotonic() - start) * 1000

            if result.get("isError"):
                content = result.get("content", [])
                if content and isinstance(content, list) and len(content) > 0:
                    error_text = content[0].get("text", "unknown error")
                else:
                    error_text = "unknown error"
                return MCPToolResult(
                    success=False,
                    error=f"MCP server error: {error_text}",
                    duration_ms=elapsed,
                )

            # MCP 1.4.0 returns ``content`` as a list of ContentItem dicts
            # (e.g. ``[{"type": "text", "text": "..."}]``).  Normalise to a
            # plain dict so that ``MCPToolResult.data`` is always
            # ``dict[str, Any]`` and downstream code can rely on dict access.
            raw_content = result.get("content", [])
            if isinstance(raw_content, list):
                normalised: dict[str, Any] = {"content": raw_content}
            elif isinstance(raw_content, dict):
                normalised = raw_content
            else:
                # Unexpected type — wrap it so the dict contract holds.
                normalised = {"content": raw_content}
            return MCPToolResult(
                success=True,
                data=normalised,
                duration_ms=elapsed,
            )

    def register_tools(
        self,
        registry: ToolRegistry,
        namespace: str,
        tool_filter: MCPToolFilter | None = None,
    ) -> list[str]:
        """Discover and register MCP tools into a ToolRegistry.

        Parameters
        ----------
        registry:
            The ``ToolRegistry`` to register tools into.
        namespace:
            Namespace prefix for registered tool names.
        tool_filter:
            Optional include/exclude filter.

        Returns
        -------
        list[str]
            Names of successfully registered tools.
        """
        descriptors = self.discover_tools(tool_filter)
        registered: list[str] = []

        for desc in descriptors:
            tool_name = f"{namespace}/{desc.name}"
            existing = registry.get(tool_name)
            if existing is not None:
                registry.remove(tool_name)

            adapter_ref = self

            def _make_handler(tn: str, ar: MCPToolAdapter) -> Any:
                def handler(**kwargs: Any) -> dict[str, Any]:
                    result = ar.invoke(tn, kwargs)
                    if not result.success:
                        return {"error": result.error}
                    return result.data

                return handler

            inferred = self.infer_capabilities(desc.name)
            # MCP tools are never checkpointable — they make real external
            # calls to a server that cannot be rolled back transactionally.
            capabilities = ToolCapability(
                read_only=inferred["read_only"],
                writes=inferred["writes"],
                checkpointable=False,
            )

            inferred_slots = self.infer_resource_slots(desc.name, desc.input_schema)
            slot_dicts: list[dict[str, str]] = [
                {
                    "name": slot.name,
                    "resource_type": slot.resource_type,
                    "access": slot.access.value,
                    "binding": slot.binding.value,
                }
                for slot in inferred_slots
            ]

            # TODO(#882): resource_slots are stored in source_metadata but
            # nothing downstream reads them yet.  The ToolRegistry persists
            # resource bindings via domain Tool.resource_slots, and the DB
            # migration (c1_001) stores them in tool_resource_bindings.
            # A follow-up ticket should wire inferred slots into the domain
            # Tool objects so the registry and DB actually consume them.
            cap_meta = self.capability_metadata
            spec = ToolSpec(
                name=tool_name,
                description=desc.description or f"MCP tool: {desc.name}",
                input_schema=desc.input_schema,
                capabilities=capabilities,
                handler=_make_handler(desc.name, adapter_ref),
                source="mcp",
                source_metadata={
                    "server": self._config.name,
                    "resource_slots": slot_dicts,
                    "annotations": desc.annotations,
                    "server_capabilities": {
                        "tools": cap_meta.tools,
                        "resources": cap_meta.resources,
                        "prompts": cap_meta.prompts,
                    },
                },
            )
            registry.register(spec)
            registered.append(tool_name)

        return registered

    # -- Parameter name patterns for resource slot inference ----------------

    #: Parameter names that indicate a file resource.
    _FILE_PARAMS: frozenset[str] = frozenset({"file_path", "filepath", "file"})

    #: Parameter names that indicate a directory resource.
    _DIR_PARAMS: frozenset[str] = frozenset({"directory", "dir", "folder"})

    #: Parameter names that indicate a git repository resource.
    _REPO_PARAMS: frozenset[str] = frozenset({"repo_path", "repository"})

    @staticmethod
    def infer_resource_slots(
        tool_name: str,
        input_schema: dict[str, Any],
    ) -> list[ResourceSlot]:
        """Infer resource slots from MCP tool parameter schemas.

        Scans ``input_schema["properties"]`` for parameter names matching
        file, directory, and repository patterns and creates
        :class:`~cleveragents.domain.models.core.tool.ResourceSlot` objects
        with appropriate metadata.

        Parameter -> ResourceSlot mapping:

        - ``file_path``, ``filepath``, ``file`` -> file / read_write / parameter
        - ``directory``, ``dir``, ``folder`` -> directory / read_only / parameter
        - ``repo_path``, ``repository`` -> repository / read_write / contextual

        Parameters
        ----------
        tool_name:
            Name of the MCP tool (used for logging).
        input_schema:
            JSON Schema dict for the tool's inputs.

        Returns
        -------
        list[ResourceSlot]
            Inferred resource slots (may be empty).
        """
        # Mapping from parameter-name sets to slot metadata.
        _SLOT_MAP: list[
            tuple[frozenset[str], str, str, ResourceAccessMode, BindingMode]
        ] = [
            (
                MCPToolAdapter._FILE_PARAMS,
                "file",
                "file",
                ResourceAccessMode.READ_WRITE,
                BindingMode.PARAMETER,
            ),
            (
                MCPToolAdapter._DIR_PARAMS,
                "directory",
                "directory",
                ResourceAccessMode.READ_ONLY,
                BindingMode.PARAMETER,
            ),
            (
                MCPToolAdapter._REPO_PARAMS,
                "repository",
                "git-checkout",
                ResourceAccessMode.READ_WRITE,
                BindingMode.CONTEXTUAL,
            ),
        ]

        properties: dict[str, Any] = input_schema.get("properties", {})
        slots: list[ResourceSlot] = []
        seen_names: set[str] = set()

        for param_name in properties:
            for param_set, slot_name, resource_type, access, binding in _SLOT_MAP:
                if param_name not in param_set:
                    continue
                if slot_name in seen_names:
                    break
                seen_names.add(slot_name)
                slots.append(
                    ResourceSlot(
                        name=slot_name,
                        resource_type=resource_type,
                        access=access,
                        binding=binding,
                        required=True,
                        description=(
                            f"{slot_name.title()} resource inferred from "
                            f"'{param_name}' parameter of MCP tool "
                            f"'{tool_name}'"
                        ),
                    )
                )
                break

        return slots

    @staticmethod
    def infer_capabilities(tool_name: str) -> dict[str, bool]:
        """Infer ToolCapability metadata from an MCP tool name.

        Applies heuristics from the specification:
        - Names containing read/get/list/search/find -> read_only
        - Names containing write/create/update/delete/set -> writes
        """
        lower = tool_name.lower()
        read_keywords = {"read", "get", "list", "search", "find"}
        write_keywords = {"write", "create", "update", "delete", "set"}

        parts = set(lower.replace("-", "_").split("_"))

        is_read = bool(parts & read_keywords)
        is_write = bool(parts & write_keywords)

        if is_read and not is_write:
            return {"read_only": True, "writes": False}
        if is_write:
            return {"read_only": False, "writes": True}
        return {"read_only": False, "writes": False}

    @staticmethod
    def _apply_filter(
        descriptors: list[MCPToolDescriptor],
        tool_filter: MCPToolFilter,
    ) -> list[MCPToolDescriptor]:
        result = descriptors
        if tool_filter.include:
            include_set = set(tool_filter.include)
            result = [d for d in result if d.name in include_set]
        if tool_filter.exclude:
            exclude_set = set(tool_filter.exclude)
            result = [d for d in result if d.name not in exclude_set]
        return result

    @staticmethod
    def _validate_input(
        descriptor: MCPToolDescriptor,
        arguments: dict[str, Any],
    ) -> str | None:
        schema = descriptor.input_schema
        if not schema:
            return None
        try:
            jsonschema.validate(instance=arguments, schema=schema)
        except jsonschema.ValidationError as exc:
            return str(exc.message)
        return None
