"""Pluggable extension interface for MCP servers (SEP-2133).

An extension is a self-contained, opt-in bundle of MCP behaviour, identified by
a reverse-DNS string (e.g. `io.modelcontextprotocol/ui`). It is passed to
`MCPServer(extensions=[...])`, and the server applies a *closed* set of
contribution kinds: tools, resources, new request methods, and one `tools/call`
interceptor. The server never hands itself to an extension; the extension
declares what it adds, and the server consumes it.

The shape follows the httpx2 `Transport`/`Auth` pattern: a narrow base class whose
methods have sensible defaults, so an extension overrides only what it needs. A
purely additive extension (Apps) overrides `tools`/`resources`; an interceptive
one overrides `methods`/`intercept_tool_call`.

This module lives at the `mcp.server` tier (not `mcp.server.mcpserver`) so the
base class itself never drags in the composition tier that consumes it;
extensions remain importable without constructing an `MCPServer`.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from mcp_types import CallToolRequestParams
from mcp_types.methods import SPEC_CLIENT_METHODS
from pydantic import BaseModel

from mcp.server.context import CallNext, HandlerResult, ServerRequestContext

# Re-exported from `mcp.shared.extension` (shared with the client surface) for existing importers.
from mcp.shared.extension import validate_extension_identifier as validate_extension_identifier

if TYPE_CHECKING:
    from mcp.server.mcpserver.resources import Resource

RequestHandler = Callable[[ServerRequestContext[Any, Any], Any], Awaitable[HandlerResult]]


@dataclass(frozen=True)
class ToolBinding:
    """A tool an extension contributes, plus the `_meta` to stamp on it."""

    fn: Callable[..., Any]
    meta: dict[str, Any] | None = None
    kwargs: dict[str, Any] = field(default_factory=lambda: {})


@dataclass(frozen=True)
class ResourceBinding:
    """A pre-built resource an extension contributes."""

    resource: Resource


@dataclass(frozen=True)
class MethodBinding:
    """A new request method an extension serves, e.g. `tasks/get`.

    `params_type` validates incoming params before `handler` runs; it should
    subclass `RequestParams` so `_meta` parses uniformly. `protocol_versions`,
    when set, restricts the method to those wire versions - a request for the
    method at any other version is rejected as `METHOD_NOT_FOUND`, mirroring the
    spec's `(method, version)` boundary table. `None` (the default) admits the
    method at every version.

    Extension methods are additive: `method` must not name a spec-defined
    request method (`tools/list`, `completion/complete`, ...) — those handlers
    belong to the server, and an extension binding one would silently shadow or
    be shadowed by it. Both constraints are enforced at construction. To
    re-provide a spec method the 2026 revision removed (e.g. `logging/setLevel`
    for legacy clients), use the lowlevel `Server.add_request_handler` API
    instead — the runner's per-version surface gate would never route such a
    method to an extension handler anyway.
    """

    method: str
    params_type: type[BaseModel]
    handler: RequestHandler
    protocol_versions: frozenset[str] | None = None

    def __post_init__(self) -> None:
        if self.method in SPEC_CLIENT_METHODS:
            raise ValueError(
                f"MethodBinding cannot bind spec method {self.method!r}; extension methods are "
                "additive — use Extension.intercept_tool_call or Server.middleware to wrap core behaviour"
            )
        if self.protocol_versions is not None and not self.protocol_versions:
            raise ValueError(
                f"MethodBinding for {self.method!r} has an empty protocol_versions set, so it could "
                "never be served; use None to admit every version"
            )


class Extension:
    """Base class for an opt-in MCP extension. Override only the methods you need.

    Subclass and set `identifier`, then override the contribution methods that
    apply. Every method has a default, so a minimal extension overrides nothing
    but `identifier` and one of `tools`/`resources`/`methods`. `identifier` is
    enforced at subclass-definition time.
    """

    #: Reverse-DNS extension identifier, advertised under `ServerCapabilities.extensions`.
    identifier: str

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        # Validate a class-level `identifier` at definition time. A subclass may
        # instead assign `identifier` in `__init__` (per-instance ids); that case
        # is validated when the extension is applied, since no class attribute
        # exists to inspect here.
        identifier = cls.__dict__.get("identifier")
        if identifier is not None:
            validate_extension_identifier(identifier, owner=cls.__name__)

    def settings(self) -> dict[str, Any]:
        """Per-extension settings advertised at `capabilities.extensions[identifier]`.

        An empty dict (the default) advertises the extension with no settings.
        """
        return {}

    def tools(self) -> Sequence[ToolBinding]:
        """Tools this extension contributes (additive)."""
        return ()

    def resources(self) -> Sequence[ResourceBinding]:
        """Resources this extension contributes (additive)."""
        return ()

    def methods(self) -> Sequence[MethodBinding]:
        """New request methods this extension serves (additive)."""
        return ()

    async def intercept_tool_call(
        self,
        params: CallToolRequestParams,
        ctx: ServerRequestContext[Any, Any],
        call_next: CallNext,
    ) -> HandlerResult:
        """Wrap `tools/call`. Default: pass through unchanged.

        Override to short-circuit (return a result without calling `call_next`)
        or to observe the call. `params` is the validated `tools/call` params;
        `call_next(ctx)` runs the rest of the chain and the real handler, and
        returns the handler's domain result. Interceptors run at the handler
        layer: whatever they return is serialized like any handler result,
        including the 2026-era `serverInfo` `_meta` stamp. The `params` this
        interceptor received is what the wrapped handler is invoked with -
        passing a rewritten context through `call_next` adjusts what the
        handler observes on `ctx`, not the tool invocation. Wire-level request
        rewriting belongs to `Server.middleware`, above params validation.
        """
        return await call_next(ctx)


def compose_tool_call_handler(extensions: Sequence[Extension], handler: RequestHandler) -> RequestHandler:
    """Fold every extension's `intercept_tool_call` around the `tools/call` handler.

    The returned handler nests the interceptors (first extension outermost) and
    replaces the plain `tools/call` registration. Interception happens at the
    handler layer, below the runner's outbound envelope pass, so a
    short-circuiting interceptor's result is sieved and stamped exactly like
    the wrapped handler's would be.
    """

    async def wrapped(ctx: ServerRequestContext[Any, Any], params: CallToolRequestParams) -> HandlerResult:
        async def innermost(inner_ctx: ServerRequestContext[Any, Any]) -> HandlerResult:
            return await handler(inner_ctx, params)

        chain: CallNext = innermost
        for extension in reversed(extensions):
            chain = _bind_interceptor(extension, params, chain)
        return await chain(ctx)

    return wrapped


def _bind_interceptor(extension: Extension, params: CallToolRequestParams, call_next: CallNext) -> CallNext:
    async def call(ctx: ServerRequestContext[Any, Any]) -> HandlerResult:
        return await extension.intercept_tool_call(params, ctx, call_next)

    return call