"""ProxyProvider for proxying to remote MCP servers.

This module provides the `ProxyProvider` class that proxies components from
a remote MCP server via a client factory. It also provides proxy component
classes that forward execution to remote servers.
"""

from __future__ import annotations

import base64
import inspect
import time
from collections.abc import Awaitable, Callable, Sequence
from typing import TYPE_CHECKING, Any, cast

import anyio
import httpx
import mcp.types
from mcp import ServerSession
from mcp.client.session import ClientSession
from mcp.server.lowlevel.server import request_ctx
from mcp.shared.context import LifespanContextT, RequestContext
from mcp.shared.exceptions import McpError
from mcp.types import (
    METHOD_NOT_FOUND,
    BlobResourceContents,
    ElicitRequestFormParams,
    TextResourceContents,
)
from pydantic.networks import AnyUrl

from fastmcp.client.client import Client, FastMCP1Server
from fastmcp.client.elicitation import ElicitResult, create_elicitation_callback
from fastmcp.client.logging import LogMessage, create_log_callback
from fastmcp.client.roots import RootsList, create_roots_callback
from fastmcp.client.sampling import create_sampling_callback
from fastmcp.client.telemetry import client_span
from fastmcp.client.transports import ClientTransportT
from fastmcp.exceptions import ResourceError
from fastmcp.mcp_config import MCPConfig
from fastmcp.prompts import Message, Prompt, PromptResult
from fastmcp.prompts.base import PromptArgument
from fastmcp.resources import Resource, ResourceTemplate
from fastmcp.resources.base import ResourceContent, ResourceResult
from fastmcp.resources.template import expand_uri_template
from fastmcp.server.context import Context
from fastmcp.server.dependencies import get_context
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp.server.providers.aggregate import ProviderErrorStrategy
from fastmcp.server.providers.base import Provider
from fastmcp.server.server import FastMCP
from fastmcp.server.tasks.config import TaskConfig
from fastmcp.tools.base import Tool, ToolResult
from fastmcp.utilities.components import FastMCPComponent, get_fastmcp_metadata
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.versions import VersionSpec, version_sort_key

if TYPE_CHECKING:
    from pathlib import Path

    from fastmcp.client.transports import ClientTransport

logger = get_logger(__name__)

# Type alias for client factory functions
ClientFactoryT = Callable[[], Client] | Callable[[], Awaitable[Client]]


def _proxy_upstream_error(error: Exception) -> McpError:
    return McpError(
        mcp.types.ErrorData(
            code=mcp.types.INTERNAL_ERROR,
            message=str(error),
        )
    )


class ProxyInitializeMiddleware(Middleware):
    def __init__(self, proxy: FastMCPProxy) -> None:
        self.proxy = proxy

    async def on_initialize(
        self,
        context: MiddlewareContext[mcp.types.InitializeRequest],
        call_next: CallNext[
            mcp.types.InitializeRequest,
            mcp.types.InitializeResult | None,
        ],
    ) -> mcp.types.InitializeResult | None:
        client = await self.proxy._get_client()
        try:
            if isinstance(client, StatefulProxyClient):
                ctx = context.fastmcp_context
                if ctx is not None:
                    client._proxy_rc_ref[0] = (
                        ctx.request_context,
                        ctx._fastmcp,
                    )
            async with client:
                await client.initialize()
        except McpError:
            raise
        except (
            RuntimeError,
            TimeoutError,
            httpx.HTTPError,
            anyio.ClosedResourceError,
            anyio.EndOfStream,
            anyio.BrokenResourceError,
        ) as error:
            raise _proxy_upstream_error(error) from error

        return await call_next(context)


# -----------------------------------------------------------------------------
# Proxy Component Classes
# -----------------------------------------------------------------------------


class ProxyTool(Tool):
    """A Tool that represents and executes a tool on a remote server."""

    task_config: TaskConfig = TaskConfig(mode="forbidden")
    _backend_name: str | None = None

    def __init__(self, client_factory: ClientFactoryT, **kwargs: Any):
        super().__init__(**kwargs)
        self._client_factory = client_factory

    async def _get_client(self) -> Client:
        """Gets a client instance by calling the sync or async factory."""
        client = self._client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    def model_copy(self, **kwargs: Any) -> ProxyTool:
        """Override to preserve _backend_name when name changes."""
        update = kwargs.get("update", {})
        if "name" in update and self._backend_name is None:
            # First time name is being changed, preserve original for backend calls
            update = {**update, "_backend_name": self.name}
            kwargs["update"] = update
        return super().model_copy(**kwargs)

    @classmethod
    def from_mcp_tool(
        cls, client_factory: ClientFactoryT, mcp_tool: mcp.types.Tool
    ) -> ProxyTool:
        """Factory method to create a ProxyTool from a raw MCP tool schema."""
        return cls(
            client_factory=client_factory,
            name=mcp_tool.name,
            title=mcp_tool.title,
            description=mcp_tool.description,
            parameters=mcp_tool.inputSchema,
            annotations=mcp_tool.annotations,
            output_schema=mcp_tool.outputSchema,
            icons=mcp_tool.icons,
            meta=mcp_tool.meta,
            tags=get_fastmcp_metadata(mcp_tool.meta).get("tags", []),
            execution=mcp_tool.execution,
        )

    async def run(
        self,
        arguments: dict[str, Any],
        context: Context | None = None,
    ) -> ToolResult:
        """Executes the tool by making a call through the client."""
        backend_name = self._backend_name or self.name
        with client_span(
            f"tools/call {backend_name}",
            "tools/call",
            backend_name,
            tool_name=backend_name,
        ) as span:
            span.set_attribute("fastmcp.provider.type", "ProxyProvider")
            client = await self._get_client()
            async with client:
                ctx = context or get_context()
                # StatefulProxyClient reuses sessions across requests, so
                # its receive-loop task has stale ContextVars from the first
                # request. Stash the current RequestContext in the shared
                # ref so handlers can restore it before forwarding.
                if isinstance(client, StatefulProxyClient):
                    client._proxy_rc_ref[0] = (
                        ctx.request_context,
                        ctx._fastmcp,  # weakref to FastMCP, not the Context
                    )
                # Build meta dict from request context
                meta: dict[str, Any] | None = None
                if hasattr(ctx, "request_context"):
                    req_ctx = ctx.request_context
                    # Start with existing meta if present
                    if hasattr(req_ctx, "meta") and req_ctx.meta:
                        meta = dict(req_ctx.meta)
                    # Add task metadata if this is a task request
                    if (
                        hasattr(req_ctx, "experimental")
                        and hasattr(req_ctx.experimental, "is_task")
                        and req_ctx.experimental.is_task
                    ):
                        task_metadata = req_ctx.experimental.task_metadata
                        if task_metadata:
                            meta = meta or {}
                            meta["modelcontextprotocol.io/task"] = (
                                task_metadata.model_dump(exclude_none=True)
                            )

                result = await client.call_tool_mcp(
                    name=backend_name, arguments=arguments, meta=meta
                )
            # Pass an upstream error result through faithfully rather than
            # collapsing it into a raised ToolError — this preserves the
            # backend's content (including non-text and structured content),
            # and the client still raises on isError by default.
            # Preserve backend's meta (includes task metadata for background tasks)
            return ToolResult(
                content=result.content,
                structured_content=result.structuredContent,
                meta=result.meta,
                is_error=result.isError,
            )

    def get_span_attributes(self) -> dict[str, Any]:
        return super().get_span_attributes() | {
            "fastmcp.provider.type": "ProxyProvider",
            "fastmcp.proxy.backend_name": self._backend_name or self.name,
        }


class ProxyResource(Resource):
    """A Resource that represents and reads a resource from a remote server."""

    task_config: TaskConfig = TaskConfig(mode="forbidden")
    _cached_content: ResourceResult | None = None
    _backend_uri: str | None = None

    def __init__(
        self,
        client_factory: ClientFactoryT,
        *,
        _cached_content: ResourceResult | None = None,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self._client_factory = client_factory
        self._cached_content = _cached_content

    async def _get_client(self) -> Client:
        """Gets a client instance by calling the sync or async factory."""
        client = self._client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    def model_copy(self, **kwargs: Any) -> ProxyResource:
        """Override to preserve _backend_uri when uri changes."""
        update = kwargs.get("update", {})
        if "uri" in update and self._backend_uri is None:
            # First time uri is being changed, preserve original for backend calls
            update = {**update, "_backend_uri": str(self.uri)}
            kwargs["update"] = update
        return super().model_copy(**kwargs)

    @classmethod
    def from_mcp_resource(
        cls,
        client_factory: ClientFactoryT,
        mcp_resource: mcp.types.Resource,
    ) -> ProxyResource:
        """Factory method to create a ProxyResource from a raw MCP resource schema."""

        return cls(
            client_factory=client_factory,
            uri=mcp_resource.uri,
            name=mcp_resource.name,
            title=mcp_resource.title,
            description=mcp_resource.description,
            mime_type=mcp_resource.mimeType or "text/plain",
            icons=mcp_resource.icons,
            meta=mcp_resource.meta,
            tags=get_fastmcp_metadata(mcp_resource.meta).get("tags", []),
            task_config=TaskConfig(mode="forbidden"),
        )

    async def read(self) -> ResourceResult:
        """Read the resource content from the remote server."""
        if self._cached_content is not None:
            return self._cached_content

        backend_uri = self._backend_uri or str(self.uri)
        with client_span(
            "resources/read",
            "resources/read",
            backend_uri,
            resource_uri=backend_uri,
        ) as span:
            span.set_attribute("fastmcp.provider.type", "ProxyProvider")
            client = await self._get_client()
            async with client:
                result = await client.read_resource(backend_uri)
            if not result:
                raise ResourceError(
                    f"Remote server returned empty content for {backend_uri}"
                )

            # Process all items in the result list, not just the first one
            contents: list[ResourceContent] = []
            for item in result:
                if isinstance(item, TextResourceContents):
                    contents.append(
                        ResourceContent(
                            content=item.text,
                            mime_type=item.mimeType,
                            meta=item.meta,
                        )
                    )
                elif isinstance(item, BlobResourceContents):
                    contents.append(
                        ResourceContent(
                            content=base64.b64decode(item.blob),
                            mime_type=item.mimeType,
                            meta=item.meta,
                        )
                    )
                else:
                    raise ResourceError(f"Unsupported content type: {type(item)}")

            return ResourceResult(contents=contents)

    def get_span_attributes(self) -> dict[str, Any]:
        return super().get_span_attributes() | {
            "fastmcp.provider.type": "ProxyProvider",
            "fastmcp.proxy.backend_uri": self._backend_uri or str(self.uri),
        }


class ProxyTemplate(ResourceTemplate):
    """A ResourceTemplate that represents and creates resources from a remote server template."""

    task_config: TaskConfig = TaskConfig(mode="forbidden")
    _backend_uri_template: str | None = None

    def __init__(self, client_factory: ClientFactoryT, **kwargs: Any):
        super().__init__(**kwargs)
        self._client_factory = client_factory

    async def _get_client(self) -> Client:
        """Gets a client instance by calling the sync or async factory."""
        client = self._client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    def model_copy(self, **kwargs: Any) -> ProxyTemplate:
        """Override to preserve _backend_uri_template when uri_template changes."""
        update = kwargs.get("update", {})
        if "uri_template" in update and self._backend_uri_template is None:
            # First time uri_template is being changed, preserve original for backend
            update = {**update, "_backend_uri_template": self.uri_template}
            kwargs["update"] = update
        return super().model_copy(**kwargs)

    @classmethod
    def from_mcp_template(  # type: ignore[override]
        cls, client_factory: ClientFactoryT, mcp_template: mcp.types.ResourceTemplate
    ) -> ProxyTemplate:  # ty:ignore[invalid-method-override]
        """Factory method to create a ProxyTemplate from a raw MCP template schema."""

        return cls(
            client_factory=client_factory,
            uri_template=mcp_template.uriTemplate,
            name=mcp_template.name,
            title=mcp_template.title,
            description=mcp_template.description,
            mime_type=mcp_template.mimeType or "text/plain",
            icons=mcp_template.icons,
            parameters={},  # Remote templates don't have local parameters
            meta=mcp_template.meta,
            tags=get_fastmcp_metadata(mcp_template.meta).get("tags", []),
            task_config=TaskConfig(mode="forbidden"),
        )

    async def create_resource(
        self,
        uri: str,
        params: dict[str, Any],
        context: Context | None = None,
    ) -> ProxyResource:
        """Create a resource from the template by calling the remote server."""
        # don't use the provided uri, because it may not be the same as the
        # uri_template on the remote server. expand_uri_template percent-encodes
        # path and query values so the backend URI round-trips correctly.
        backend_template = self._backend_uri_template or self.uri_template
        parameterized_uri = expand_uri_template(backend_template, params)
        client = await self._get_client()
        async with client:
            result = await client.read_resource(parameterized_uri)

        if not result:
            raise ResourceError(
                f"Remote server returned empty content for {parameterized_uri}"
            )

        # Process all items in the result list, not just the first one
        contents: list[ResourceContent] = []
        for item in result:
            if isinstance(item, TextResourceContents):
                contents.append(
                    ResourceContent(
                        content=item.text,
                        mime_type=item.mimeType,
                        meta=item.meta,
                    )
                )
            elif isinstance(item, BlobResourceContents):
                contents.append(
                    ResourceContent(
                        content=base64.b64decode(item.blob),
                        mime_type=item.mimeType,
                        meta=item.meta,
                    )
                )
            else:
                raise ResourceError(f"Unsupported content type: {type(item)}")

        cached_content = ResourceResult(contents=contents)

        return ProxyResource(
            client_factory=self._client_factory,
            uri=parameterized_uri,
            name=self.name,
            title=self.title,
            description=self.description,
            mime_type=result[
                0
            ].mimeType,  # Use first item's mimeType for backward compatibility
            icons=self.icons,
            meta=self.meta,
            tags=get_fastmcp_metadata(self.meta).get("tags", []),
            _cached_content=cached_content,
        )

    def get_span_attributes(self) -> dict[str, Any]:
        return super().get_span_attributes() | {
            "fastmcp.provider.type": "ProxyProvider",
            "fastmcp.proxy.backend_uri_template": (
                self._backend_uri_template or self.uri_template
            ),
        }


class ProxyPrompt(Prompt):
    """A Prompt that represents and renders a prompt from a remote server."""

    task_config: TaskConfig = TaskConfig(mode="forbidden")
    _backend_name: str | None = None

    def __init__(self, client_factory: ClientFactoryT, **kwargs):
        super().__init__(**kwargs)
        self._client_factory = client_factory

    async def _get_client(self) -> Client:
        """Gets a client instance by calling the sync or async factory."""
        client = self._client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    def model_copy(self, **kwargs: Any) -> ProxyPrompt:
        """Override to preserve _backend_name when name changes."""
        update = kwargs.get("update", {})
        if "name" in update and self._backend_name is None:
            # First time name is being changed, preserve original for backend calls
            update = {**update, "_backend_name": self.name}
            kwargs["update"] = update
        return super().model_copy(**kwargs)

    @classmethod
    def from_mcp_prompt(
        cls, client_factory: ClientFactoryT, mcp_prompt: mcp.types.Prompt
    ) -> ProxyPrompt:
        """Factory method to create a ProxyPrompt from a raw MCP prompt schema."""
        arguments = [
            PromptArgument(
                name=arg.name,
                description=arg.description,
                required=arg.required or False,
            )
            for arg in mcp_prompt.arguments or []
        ]
        return cls(
            client_factory=client_factory,
            name=mcp_prompt.name,
            title=mcp_prompt.title,
            description=mcp_prompt.description,
            arguments=arguments,
            icons=mcp_prompt.icons,
            meta=mcp_prompt.meta,
            tags=get_fastmcp_metadata(mcp_prompt.meta).get("tags", []),
            task_config=TaskConfig(mode="forbidden"),
        )

    async def render(self, arguments: dict[str, Any]) -> PromptResult:  # type: ignore[override]  # ty:ignore[invalid-method-override]
        """Render the prompt by making a call through the client."""
        backend_name = self._backend_name or self.name
        with client_span(
            f"prompts/get {backend_name}",
            "prompts/get",
            backend_name,
            prompt_name=backend_name,
        ) as span:
            span.set_attribute("fastmcp.provider.type", "ProxyProvider")
            client = await self._get_client()
            async with client:
                result = await client.get_prompt(backend_name, arguments)
            # Convert GetPromptResult to PromptResult, preserving meta from result
            # (not the static prompt meta which includes fastmcp tags)
            # Convert PromptMessages to Messages
            messages = [
                Message(content=m.content, role=m.role) for m in result.messages
            ]
            return PromptResult(
                messages=messages,
                description=result.description,
                meta=result.meta,
            )

    def get_span_attributes(self) -> dict[str, Any]:
        return super().get_span_attributes() | {
            "fastmcp.provider.type": "ProxyProvider",
            "fastmcp.proxy.backend_name": self._backend_name or self.name,
        }


# -----------------------------------------------------------------------------
# ProxyProvider
# -----------------------------------------------------------------------------


class _CacheEntry:
    """A cached sequence of components with a monotonic timestamp."""

    __slots__ = ("items", "timestamp")

    def __init__(self, items: Sequence[Any], timestamp: float):
        self.items = items
        self.timestamp = timestamp

    def is_fresh(self, ttl: float) -> bool:
        return (time.monotonic() - self.timestamp) < ttl


_DEFAULT_CACHE_TTL: float = 300.0


class ProxyProvider(Provider):
    """Provider that proxies to a remote MCP server via a client factory.

    This provider fetches components from a remote server and returns Proxy*
    component instances that forward execution to the remote server.

    All components returned by this provider have task_config.mode="forbidden"
    because tasks cannot be executed through a proxy.

    Component lists (tools, resources, templates, prompts) are cached so that
    individual lookups (e.g. during ``call_tool``) can resolve from the cache
    instead of opening a new backend connection.  The cache stores the
    backend's raw component metadata and is shared across all sessions;
    per-session visibility and auth filtering are applied after cache lookup
    by the server layer.  The cache is refreshed whenever a ``list_*`` call
    is made, and entries expire after ``cache_ttl`` seconds (default 300).
    Set ``cache_ttl=0`` to disable caching.  Disabling is recommended for
    backends whose component lists change dynamically.

    Example:
        ```python
        from fastmcp import FastMCP
        from fastmcp.server.providers.proxy import ProxyProvider, ProxyClient

        # Create a proxy provider for a remote server
        proxy = ProxyProvider(lambda: ProxyClient("http://localhost:8000/mcp"))

        mcp = FastMCP("Proxy Server")
        mcp.add_provider(proxy)

        # Can also add with namespace
        mcp.add_provider(proxy.with_namespace("remote"))
        ```
    """

    def __init__(
        self,
        client_factory: ClientFactoryT,
        cache_ttl: float | None = None,
    ):
        """Initialize a ProxyProvider.

        Args:
            client_factory: A callable that returns a Client instance when called.
                           This gives you full control over session creation and reuse.
                           Can be either a synchronous or asynchronous function.
            cache_ttl: How long (in seconds) to cache component lists for
                      individual lookups.  Defaults to 300.  Set to 0 to
                      disable caching.
        """
        super().__init__()
        self.client_factory = client_factory
        self._cache_ttl = cache_ttl if cache_ttl is not None else _DEFAULT_CACHE_TTL
        self._tools_cache: _CacheEntry[Tool] | None = None
        self._resources_cache: _CacheEntry[Resource] | None = None
        self._templates_cache: _CacheEntry[ResourceTemplate] | None = None
        self._prompts_cache: _CacheEntry[Prompt] | None = None

    async def _get_client(self) -> Client:
        """Gets a client instance by calling the sync or async factory."""
        client = self.client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    # -------------------------------------------------------------------------
    # Tool methods
    # -------------------------------------------------------------------------

    async def _list_tools(self) -> Sequence[Tool]:
        """List all tools from the remote server."""
        try:
            client = await self._get_client()
            async with client:
                mcp_tools = await client.list_tools()
                tools = [
                    ProxyTool.from_mcp_tool(self.client_factory, t) for t in mcp_tools
                ]
        except McpError as e:
            if e.error.code == METHOD_NOT_FOUND:
                tools = []
            else:
                raise
        self._tools_cache = _CacheEntry(tools, time.monotonic())
        return tools

    async def _get_tool(
        self, name: str, version: VersionSpec | None = None
    ) -> Tool | None:
        cache = self._tools_cache
        if cache is None or not cache.is_fresh(self._cache_ttl):
            await self._list_tools()
            cache = self._tools_cache
        assert cache is not None
        matching = [t for t in cache.items if t.name == name]
        if version:
            matching = [t for t in matching if version.matches(t.version)]
        if not matching:
            return None
        return max(matching, key=version_sort_key)

    # -------------------------------------------------------------------------
    # Resource methods
    # -------------------------------------------------------------------------

    async def _list_resources(self) -> Sequence[Resource]:
        """List all resources from the remote server."""
        try:
            client = await self._get_client()
            async with client:
                mcp_resources = await client.list_resources()
                resources = [
                    ProxyResource.from_mcp_resource(self.client_factory, r)
                    for r in mcp_resources
                ]
        except McpError as e:
            if e.error.code == METHOD_NOT_FOUND:
                resources = []
            else:
                raise
        self._resources_cache = _CacheEntry(resources, time.monotonic())
        return resources

    async def _get_resource(
        self, uri: str, version: VersionSpec | None = None
    ) -> Resource | None:
        cache = self._resources_cache
        if cache is None or not cache.is_fresh(self._cache_ttl):
            await self._list_resources()
            cache = self._resources_cache
        assert cache is not None
        matching = [r for r in cache.items if str(r.uri) == uri]
        if version:
            matching = [r for r in matching if version.matches(r.version)]
        if not matching:
            return None
        return max(matching, key=version_sort_key)

    # -------------------------------------------------------------------------
    # Resource template methods
    # -------------------------------------------------------------------------

    async def _list_resource_templates(self) -> Sequence[ResourceTemplate]:
        """List all resource templates from the remote server."""
        try:
            client = await self._get_client()
            async with client:
                mcp_templates = await client.list_resource_templates()
                templates = [
                    ProxyTemplate.from_mcp_template(self.client_factory, t)
                    for t in mcp_templates
                ]
        except McpError as e:
            if e.error.code == METHOD_NOT_FOUND:
                templates = []
            else:
                raise
        self._templates_cache = _CacheEntry(templates, time.monotonic())
        return templates

    async def _get_resource_template(
        self, uri: str, version: VersionSpec | None = None
    ) -> ResourceTemplate | None:
        cache = self._templates_cache
        if cache is None or not cache.is_fresh(self._cache_ttl):
            await self._list_resource_templates()
            cache = self._templates_cache
        assert cache is not None
        matching = [t for t in cache.items if t.matches(uri) is not None]
        if version:
            matching = [t for t in matching if version.matches(t.version)]
        if not matching:
            return None
        return max(matching, key=version_sort_key)

    # -------------------------------------------------------------------------
    # Prompt methods
    # -------------------------------------------------------------------------

    async def _list_prompts(self) -> Sequence[Prompt]:
        """List all prompts from the remote server."""
        try:
            client = await self._get_client()
            async with client:
                mcp_prompts = await client.list_prompts()
                prompts = [
                    ProxyPrompt.from_mcp_prompt(self.client_factory, p)
                    for p in mcp_prompts
                ]
        except McpError as e:
            if e.error.code == METHOD_NOT_FOUND:
                prompts = []
            else:
                raise
        self._prompts_cache = _CacheEntry(prompts, time.monotonic())
        return prompts

    async def _get_prompt(
        self, name: str, version: VersionSpec | None = None
    ) -> Prompt | None:
        cache = self._prompts_cache
        if cache is None or not cache.is_fresh(self._cache_ttl):
            await self._list_prompts()
            cache = self._prompts_cache
        assert cache is not None
        matching = [p for p in cache.items if p.name == name]
        if version:
            matching = [p for p in matching if version.matches(p.version)]
        if not matching:
            return None
        return max(matching, key=version_sort_key)

    # -------------------------------------------------------------------------
    # Task methods
    # -------------------------------------------------------------------------

    async def get_tasks(self) -> Sequence[FastMCPComponent]:
        """Return empty list since proxy components don't support tasks.

        Override the base implementation to avoid calling list_tools() during
        server lifespan initialization, which would open the client before any
        context is set. All Proxy* components have task_config.mode="forbidden".
        """
        return []

    # lifespan() uses default implementation (empty context manager)
    # because client cleanup is handled per-request


# -----------------------------------------------------------------------------
# Factory Functions
# -----------------------------------------------------------------------------


def _create_client_factory(
    target: (
        Client[ClientTransportT]
        | ClientTransport
        | FastMCP[Any]
        | FastMCP1Server
        | AnyUrl
        | Path
        | MCPConfig
        | dict[str, Any]
        | str
    ),
) -> ClientFactoryT:
    """Create a client factory from the given target.

    Internal helper that handles the session strategy based on the target type:
    - Connected Client: reuses existing session (with warning about context mixing)
    - Disconnected Client: creates fresh sessions per request
    - Other targets: creates ProxyClient and fresh sessions per request
    """
    if isinstance(target, Client):
        client = target

        # Plain Clients used as proxy backends also need header forwarding,
        # same as ProxyClient (which sets this in __init__).
        from fastmcp.client.transports.http import StreamableHttpTransport
        from fastmcp.client.transports.sse import SSETransport

        if isinstance(client.transport, StreamableHttpTransport | SSETransport):
            client.transport.forward_incoming_headers = True

        if client.is_connected() and type(client) is ProxyClient:
            logger.info(
                "Proxy detected connected ProxyClient - creating fresh sessions for each "
                "request to avoid request context leakage."
            )

            def fresh_client_factory() -> Client:
                return client.new()

            return fresh_client_factory

        if client.is_connected():
            logger.info(
                "Proxy detected connected client - reusing existing session for all requests. "
                "This may cause context mixing in concurrent scenarios."
            )

            def reuse_client_factory() -> Client:
                return client

            return reuse_client_factory

        def fresh_client_factory() -> Client:
            return client.new()

        return fresh_client_factory
    else:
        # target is not a Client, so it's compatible with ProxyClient.__init__
        base_client = ProxyClient(cast(Any, target))

        def proxy_client_factory() -> Client:
            return base_client.new()

        return proxy_client_factory


# -----------------------------------------------------------------------------
# FastMCPProxy - Convenience Wrapper
# -----------------------------------------------------------------------------


class FastMCPProxy(FastMCP):
    """A FastMCP server that acts as a proxy to a remote MCP-compliant server.

    This is a convenience wrapper that creates a FastMCP server with a
    ProxyProvider. For more control, use FastMCP with add_provider(ProxyProvider(...)).

    Example:
        ```python
        from fastmcp.server import create_proxy
        from fastmcp.server.providers.proxy import FastMCPProxy, ProxyClient

        # Create a proxy server using create_proxy (recommended)
        proxy = create_proxy("http://localhost:8000/mcp")

        # Or use FastMCPProxy directly with explicit client factory
        proxy = FastMCPProxy(client_factory=lambda: ProxyClient("http://localhost:8000/mcp"))
        ```
    """

    def __init__(
        self,
        *,
        client_factory: ClientFactoryT,
        provider_error_strategy: ProviderErrorStrategy = "warn",
        **kwargs,
    ):
        """Initialize the proxy server.

        FastMCPProxy requires explicit session management via client_factory.
        Use create_proxy() for convenience with automatic session strategy.

        Args:
            client_factory: A callable that returns a Client instance when called.
                           This gives you full control over session creation and reuse.
                           Can be either a synchronous or asynchronous function.
            provider_error_strategy: How provider errors should affect aggregate
                operations. Defaults to ``"warn"`` for compatibility; use
                ``"raise"`` when the proxy should surface upstream failures.
            **kwargs: Additional settings for the FastMCP server.
        """
        super().__init__(**kwargs)
        self.provider_error_strategy = provider_error_strategy
        self.client_factory = client_factory
        provider: Provider = ProxyProvider(client_factory)
        self.add_provider(provider)
        self.middleware.append(ProxyInitializeMiddleware(self))
        self._setup_proxy_ping_handler()

    async def _get_client(self) -> Client:
        client = self.client_factory()
        if inspect.isawaitable(client):
            client = cast(Client, await client)
        return client

    def _setup_proxy_ping_handler(self) -> None:
        async def ping_remote(
            _request: mcp.types.PingRequest,
        ) -> mcp.types.ServerResult:
            client = await self._get_client()
            async with client:
                await client.ping()
            return mcp.types.ServerResult(mcp.types.EmptyResult())

        self._mcp_server.request_handlers[mcp.types.PingRequest] = ping_remote


# -----------------------------------------------------------------------------
# ProxyClient and Related
# -----------------------------------------------------------------------------


async def default_proxy_roots_handler(
    context: RequestContext[ClientSession, LifespanContextT],
) -> RootsList:
    """Forward list roots request from remote server to proxy's connected clients."""
    ctx = get_context()
    return await ctx.list_roots()


async def default_proxy_sampling_handler(
    messages: list[mcp.types.SamplingMessage],
    params: mcp.types.CreateMessageRequestParams,
    context: RequestContext[ClientSession, LifespanContextT],
) -> mcp.types.CreateMessageResult:
    """Forward sampling request from remote server to proxy's connected clients."""
    ctx = get_context()
    result = await ctx.sample(
        list(messages),
        system_prompt=params.systemPrompt,
        temperature=params.temperature,
        max_tokens=params.maxTokens,
        model_preferences=params.modelPreferences,
    )
    content = mcp.types.TextContent(type="text", text=result.text or "")
    return mcp.types.CreateMessageResult(
        role="assistant",
        model="fastmcp-client",
        # TODO(ty): remove when ty supports isinstance exclusion narrowing
        content=content,
    )


async def default_proxy_elicitation_handler(
    message: str,
    response_type: type,
    params: mcp.types.ElicitRequestParams,
    context: RequestContext[ClientSession, LifespanContextT],
) -> ElicitResult:
    """Forward elicitation request from remote server to proxy's connected clients."""
    ctx = get_context()
    # requestedSchema only exists on ElicitRequestFormParams, not ElicitRequestURLParams
    requested_schema = (
        params.requestedSchema
        if isinstance(params, ElicitRequestFormParams)
        else {"type": "object", "properties": {}}
    )
    result = await ctx.session.elicit(
        message=message,
        requestedSchema=requested_schema,
        related_request_id=ctx.request_id,
    )
    return ElicitResult(action=result.action, content=result.content)


async def default_proxy_log_handler(message: LogMessage) -> None:
    """Forward log notification from remote server to proxy's connected clients."""
    ctx = get_context()
    msg = message.data.get("msg")
    extra = message.data.get("extra")
    await ctx.log(msg, level=message.level, logger_name=message.logger, extra=extra)


async def default_proxy_progress_handler(
    progress: float,
    total: float | None,
    message: str | None,
) -> None:
    """Forward progress notification from remote server to proxy's connected clients."""
    ctx = get_context()
    await ctx.report_progress(progress, total, message)


def _restore_request_context(
    rc_ref: list[Any],
) -> None:
    """Set the ``request_ctx``, ``_current_context`` and ``_current_server``
    ContextVars from stashed values.

    Called at the start of proxy handler invocations in
    ``StatefulProxyClient`` to fix stale ContextVars in the receive-loop
    task.  Only overrides when the ContextVar is genuinely stale (same
    session, different request_id) to avoid corrupting the concurrent
    case where multiple sessions share the same ref via ``copy.copy``.

    We stash a ``(RequestContext, weakref[FastMCP])`` tuple — never a
    ``Context`` instance — because ``Context`` properties are themselves
    ContextVar-dependent and would resolve stale values in the receive
    loop.  Instead we construct a fresh ``Context`` here after restoring
    ``request_ctx``, so its property accesses read the correct values.

    This is a set-only repair of a long-lived task's ContextVars, not a
    scope: we never ``reset()`` because the prior values are stale and
    the loop keeps running.  ``_current_server`` is restored alongside
    ``_current_context`` so handlers that resolve the server via
    dependency injection (e.g. ``get_server()``) see the right instance;
    it is set directly rather than via ``Context.__aenter__`` to avoid
    opening a context-manager lifecycle on an unscoped path.
    """
    import weakref

    from fastmcp.server.context import Context, _current_context
    from fastmcp.server.dependencies import _current_server

    stashed = rc_ref[0]
    if stashed is None:
        return

    rc, fastmcp_ref = stashed
    try:
        current_rc = request_ctx.get()
    except LookupError:
        request_ctx.set(rc)
        fastmcp = fastmcp_ref()
        if fastmcp is not None:
            _current_context.set(Context(fastmcp))
            _current_server.set(weakref.ref(fastmcp))
        return
    if current_rc.session is rc.session and current_rc.request_id != rc.request_id:
        request_ctx.set(rc)
        fastmcp = fastmcp_ref()
        if fastmcp is not None:
            _current_context.set(Context(fastmcp))
            _current_server.set(weakref.ref(fastmcp))


def _make_restoring_handler(handler: Callable, rc_ref: list[Any]) -> Callable:
    """Wrap a proxy handler to restore request_ctx before delegating.

    The wrapper is a plain ``async def`` so it passes
    ``inspect.isfunction()`` checks in handler registration paths
    (e.g., ``create_roots_callback``).
    """

    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        _restore_request_context(rc_ref)
        return await handler(*args, **kwargs)

    return wrapper


class ProxyClient(Client[ClientTransportT]):
    """A proxy client that forwards advanced interactions between a remote MCP server and the proxy's connected clients.

    Supports forwarding roots, sampling, elicitation, logging, and progress.
    """

    def __init__(
        self,
        transport: ClientTransportT
        | FastMCP[Any]
        | FastMCP1Server
        | AnyUrl
        | Path
        | MCPConfig
        | dict[str, Any]
        | str,
        **kwargs,
    ):
        if "name" not in kwargs:
            kwargs["name"] = self.generate_name()
        if "roots" not in kwargs:
            kwargs["roots"] = default_proxy_roots_handler
        if "sampling_handler" not in kwargs:
            kwargs["sampling_handler"] = default_proxy_sampling_handler
        if "elicitation_handler" not in kwargs:
            kwargs["elicitation_handler"] = default_proxy_elicitation_handler
        if "log_handler" not in kwargs:
            kwargs["log_handler"] = default_proxy_log_handler
        if "progress_handler" not in kwargs:
            kwargs["progress_handler"] = default_proxy_progress_handler
        super().__init__(transport=transport, **kwargs)  # ty: ignore[no-matching-overload]

        # Enable forwarding of inbound HTTP headers (e.g. authorization) to
        # the upstream server. This is only appropriate for proxy clients,
        # where the caller's credentials should be propagated.
        from fastmcp.client.transports.http import StreamableHttpTransport
        from fastmcp.client.transports.sse import SSETransport

        if isinstance(self.transport, StreamableHttpTransport | SSETransport):
            self.transport.forward_incoming_headers = True


class StatefulProxyClient(ProxyClient[ClientTransportT]):
    """A proxy client that provides a stateful client factory for the proxy server.

    The stateful proxy client bound its copy to the server session.
    And it will be disconnected when the session is exited.

    This is useful to proxy a stateful mcp server such as the Playwright MCP server.
    Note that it is essential to ensure that the proxy server itself is also stateful.

    Because session reuse means the receive-loop task inherits a stale
    ``request_ctx`` ContextVar snapshot, the default proxy handlers are
    replaced with versions that restore the ContextVar before forwarding.
    ``ProxyTool.run`` stashes the current ``RequestContext`` in
    ``_proxy_rc_ref`` before each backend call, and the handlers consult
    it to detect (and correct) staleness.
    """

    # Mutable list shared across copies (Client.new() uses copy.copy,
    # which preserves references to mutable containers).  ProxyTool.run
    # writes [0] before each backend call; handlers read it to detect
    # stale ContextVars and restore the correct request_ctx.
    #
    # Stores a (RequestContext, weakref[FastMCP]) tuple — never a Context
    # instance — because Context properties are ContextVar-dependent and
    # would resolve stale values in the receive loop.  The restore helper
    # constructs a fresh Context from the weakref after setting request_ctx.
    _proxy_rc_ref: list[Any]
    _proxy_restoring_handler_keys: set[str]

    def __init__(self, *args: Any, **kwargs: Any):
        # Install context-restoring handler wrappers BEFORE super().__init__
        # registers them with the Client's session kwargs.
        self._proxy_rc_ref = [None]
        self._proxy_restoring_handler_keys = set()
        for key, default_fn in (
            ("roots", default_proxy_roots_handler),
            ("sampling_handler", default_proxy_sampling_handler),
            ("elicitation_handler", default_proxy_elicitation_handler),
            ("log_handler", default_proxy_log_handler),
            ("progress_handler", default_proxy_progress_handler),
        ):
            if key not in kwargs:
                kwargs[key] = _make_restoring_handler(default_fn, self._proxy_rc_ref)
                self._proxy_restoring_handler_keys.add(key)

        super().__init__(*args, **kwargs)
        self._caches: dict[ServerSession, Client[ClientTransportT]] = {}

    def _bind_restoring_handlers(self) -> None:
        if "roots" in self._proxy_restoring_handler_keys:
            self._session_kwargs["list_roots_callback"] = create_roots_callback(
                _make_restoring_handler(default_proxy_roots_handler, self._proxy_rc_ref)
            )
        if "sampling_handler" in self._proxy_restoring_handler_keys:
            self._session_kwargs["sampling_callback"] = create_sampling_callback(
                _make_restoring_handler(
                    default_proxy_sampling_handler, self._proxy_rc_ref
                )
            )
        if "elicitation_handler" in self._proxy_restoring_handler_keys:
            self._session_kwargs["elicitation_callback"] = create_elicitation_callback(
                _make_restoring_handler(
                    default_proxy_elicitation_handler, self._proxy_rc_ref
                )
            )
        if "log_handler" in self._proxy_restoring_handler_keys:
            self._session_kwargs["logging_callback"] = create_log_callback(
                _make_restoring_handler(default_proxy_log_handler, self._proxy_rc_ref)
            )
        if "progress_handler" in self._proxy_restoring_handler_keys:
            self._progress_handler = _make_restoring_handler(
                default_proxy_progress_handler, self._proxy_rc_ref
            )

    def new(self) -> StatefulProxyClient[ClientTransportT]:
        new_client = cast(StatefulProxyClient[ClientTransportT], super().new())
        new_client._proxy_rc_ref = [None]
        new_client._bind_restoring_handlers()
        return new_client

    async def __aexit__(self, exc_type, exc_value, traceback) -> None:  # type: ignore[override]  # ty:ignore[invalid-method-override]
        """The stateful proxy client will be forced disconnected when the session is exited.

        So we do nothing here.
        """

    async def clear(self):
        """Clear all cached clients and force disconnect them."""
        while self._caches:
            _, cache = self._caches.popitem()
            await cache._disconnect(force=True)

    def new_stateful(self) -> Client[ClientTransportT]:
        """Create a new stateful proxy client instance with the same configuration.

        Use this method as the client factory for stateful proxy server.
        """
        session = get_context().session
        proxy_client = self._caches.get(session, None)

        if proxy_client is None:
            proxy_client = self.new()
            logger.debug(f"{proxy_client} created for {session}")
            self._caches[session] = proxy_client

            async def _on_session_exit():
                self._caches.pop(session, None)
                logger.debug(f"{proxy_client} will be disconnect")
                # This callback runs while the server session's exit stack is
                # unwinding, which usually happens because the owning task is
                # being cancelled. Shield the disconnect so the forced cleanup
                # actually runs to completion instead of aborting at the first
                # cancellation checkpoint (e.g. acquiring the session lock).
                with anyio.CancelScope(shield=True):
                    await proxy_client._disconnect(force=True)

            session._exit_stack.push_async_callback(_on_session_exit)

        return proxy_client