"""FastMCP - A more ergonomic interface for MCP servers."""

from __future__ import annotations as _annotations

import inspect
import re
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
from contextlib import (
    AbstractAsyncContextManager,
    asynccontextmanager,
)
from itertools import chain
from typing import Any, Generic, Literal

import anyio
import pydantic_core
from pydantic import BaseModel, Field
from pydantic.networks import AnyUrl
from pydantic_settings import BaseSettings, SettingsConfigDict
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.authentication import AuthenticationMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount, Route
from starlette.types import Receive, Scope, Send

from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
from mcp.server.auth.middleware.bearer_auth import (
    BearerAuthBackend,
    RequireAuthMiddleware,
)
from mcp.server.auth.provider import OAuthAuthorizationServerProvider
from mcp.server.auth.settings import (
    AuthSettings,
)
from mcp.server.fastmcp.exceptions import ResourceError
from mcp.server.fastmcp.prompts import Prompt, PromptManager
from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
from mcp.server.fastmcp.tools import ToolManager
from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger
from mcp.server.fastmcp.utilities.types import Image
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server.lowlevel.server import LifespanResultT
from mcp.server.lowlevel.server import Server as MCPServer
from mcp.server.lowlevel.server import lifespan as default_lifespan
from mcp.server.session import ServerSession, ServerSessionT
from mcp.server.sse import SseServerTransport
from mcp.server.stdio import stdio_server
from mcp.server.streamable_http import EventStore

# Try to import ETDI components
# Lazy import ETDI to avoid circular dependencies
ETDI_AVAILABLE = False
try:
    import mcp.etdi.types
    ETDI_AVAILABLE = True
except ImportError:
    pass


class SecurityError(Exception):
    """ETDI security violation"""
    pass
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.shared.context import LifespanContextT, RequestContext
from mcp.types import (
    AnyFunction,
    EmbeddedResource,
    GetPromptResult,
    ImageContent,
    TextContent,
    ToolAnnotations,
)
from mcp.types import Prompt as MCPPrompt
from mcp.types import PromptArgument as MCPPromptArgument
from mcp.types import Resource as MCPResource
from mcp.types import ResourceTemplate as MCPResourceTemplate
from mcp.types import Tool as MCPTool

logger = get_logger(__name__)


class Settings(BaseSettings, Generic[LifespanResultT]):
    """FastMCP server settings.

    All settings can be configured via environment variables with the prefix FASTMCP_.
    For example, FASTMCP_DEBUG=true will set debug=True.
    """

    model_config = SettingsConfigDict(
        env_prefix="FASTMCP_",
        env_file=".env",
        env_nested_delimiter="__",
        nested_model_default_partial_update=True,
        extra="ignore",
    )

    # Server settings
    debug: bool = False
    log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"

    # HTTP settings
    host: str = "0.0.0.0"
    port: int = 8000
    mount_path: str = "/"  # Mount path (e.g. "/github", defaults to root path)
    sse_path: str = "/sse"
    message_path: str = "/messages/"
    streamable_http_path: str = "/mcp"

    # StreamableHTTP settings
    json_response: bool = False
    stateless_http: bool = (
        False  # If True, uses true stateless mode (new transport per request)
    )

    # resource settings
    warn_on_duplicate_resources: bool = True

    # tool settings
    warn_on_duplicate_tools: bool = True

    # prompt settings
    warn_on_duplicate_prompts: bool = True

    dependencies: list[str] = Field(
        default_factory=list,
        description="List of dependencies to install in the server environment",
    )

    lifespan: (
        Callable[[FastMCP], AbstractAsyncContextManager[LifespanResultT]] | None
    ) = Field(None, description="Lifespan context manager")

    auth: AuthSettings | None = None


def lifespan_wrapper(
    app: FastMCP,
    lifespan: Callable[[FastMCP], AbstractAsyncContextManager[LifespanResultT]],
) -> Callable[[MCPServer[LifespanResultT]], AbstractAsyncContextManager[object]]:
    @asynccontextmanager
    async def wrap(s: MCPServer[LifespanResultT]) -> AsyncIterator[object]:
        async with lifespan(app) as context:
            yield context

    return wrap


class FastMCP:
    def __init__(
        self,
        name: str | None = None,
        instructions: str | None = None,
        auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any]
        | None = None,
        event_store: EventStore | None = None,
        **settings: Any,
    ):
        self.settings = Settings(**settings)

        self._mcp_server = MCPServer(
            name=name or "FastMCP",
            instructions=instructions,
            lifespan=(
                lifespan_wrapper(self, self.settings.lifespan)
                if self.settings.lifespan
                else default_lifespan
            ),
        )
        self._tool_manager = ToolManager(
            warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools
        )
        self._resource_manager = ResourceManager(
            warn_on_duplicate_resources=self.settings.warn_on_duplicate_resources
        )
        self._prompt_manager = PromptManager(
            warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts
        )
        if (self.settings.auth is not None) != (auth_server_provider is not None):
            # TODO: after we support separate authorization servers (see
            # https://github.com/modelcontextprotocol/modelcontextprotocol/pull/284)
            # we should validate that if auth is enabled, we have either an
            # auth_server_provider to host our own authorization server,
            # OR the URL of a 3rd party authorization server.
            raise ValueError(
                "settings.auth must be specified if and only if auth_server_provider "
                "is specified"
            )
        self._auth_server_provider = auth_server_provider
        self._event_store = event_store
        self._custom_starlette_routes: list[Route] = []
        self.dependencies = self.settings.dependencies
        self._session_manager: StreamableHTTPSessionManager | None = None
        
        # ETDI security components
        if ETDI_AVAILABLE:
            try:
                from mcp.etdi import CallStackVerifier
                self._etdi_verifier = CallStackVerifier()
            except ImportError:
                self._etdi_verifier = None
            self._current_session_id = "fastmcp_session"
            self._current_user_permissions = []  # Will be set by auth middleware
        else:
            self._etdi_verifier = None

        # Set up MCP protocol handlers
        self._setup_handlers()

        # Configure logging
        configure_logging(self.settings.log_level)

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

    @property
    def instructions(self) -> str | None:
        return self._mcp_server.instructions

    @property
    def session_manager(self) -> StreamableHTTPSessionManager:
        """Get the StreamableHTTP session manager.

        This is exposed to enable advanced use cases like mounting multiple
        FastMCP servers in a single FastAPI application.

        Raises:
            RuntimeError: If called before streamable_http_app() has been called.
        """
        if self._session_manager is None:
            raise RuntimeError(
                "Session manager can only be accessed after"
                "calling streamable_http_app()."
                "The session manager is created lazily"
                "to avoid unnecessary initialization."
            )
        return self._session_manager

    def run(
        self,
        transport: Literal["stdio", "sse", "streamable-http"] = "stdio",
        mount_path: str | None = None,
    ) -> None:
        """Run the FastMCP server. Note this is a synchronous function.

        Args:
            transport: Transport protocol to use ("stdio", "sse", or "streamable-http")
            mount_path: Optional mount path for SSE transport
        """
        TRANSPORTS = Literal["stdio", "sse", "streamable-http"]
        if transport not in TRANSPORTS.__args__:  # type: ignore
            raise ValueError(f"Unknown transport: {transport}")

        match transport:
            case "stdio":
                anyio.run(self.run_stdio_async)
            case "sse":
                anyio.run(lambda: self.run_sse_async(mount_path))
            case "streamable-http":
                anyio.run(self.run_streamable_http_async)

    def _setup_handlers(self) -> None:
        """Set up core MCP protocol handlers."""
        self._mcp_server.list_tools()(self.list_tools)
        self._mcp_server.call_tool()(self.call_tool)
        self._mcp_server.list_resources()(self.list_resources)
        self._mcp_server.read_resource()(self.read_resource)
        self._mcp_server.list_prompts()(self.list_prompts)
        self._mcp_server.get_prompt()(self.get_prompt)
        self._mcp_server.list_resource_templates()(self.list_resource_templates)

    async def list_tools(self) -> list[MCPTool]:
        """List all available tools."""
        tools = self._tool_manager.list_tools()
        return [
            MCPTool(
                name=info.name,
                description=info.description,
                inputSchema=info.parameters,
                annotations=info.annotations,
            )
            for info in tools
        ]

    def get_context(self) -> Context[ServerSession, object]:
        """
        Returns a Context object. Note that the context will only be valid
        during a request; outside a request, most methods will error.
        """
        try:
            request_context = self._mcp_server.request_context
        except LookupError:
            request_context = None
        return Context(request_context=request_context, fastmcp=self)

    async def call_tool(
        self, name: str, arguments: dict[str, Any]
    ) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
        """Call a tool by name with arguments."""
        context = self.get_context()
        result = await self._tool_manager.call_tool(name, arguments, context=context)
        converted_result = _convert_to_content(result)
        return converted_result

    async def list_resources(self) -> list[MCPResource]:
        """List all available resources."""

        resources = self._resource_manager.list_resources()
        return [
            MCPResource(
                uri=resource.uri,
                name=resource.name or "",
                description=resource.description,
                mimeType=resource.mime_type,
            )
            for resource in resources
        ]

    async def list_resource_templates(self) -> list[MCPResourceTemplate]:
        templates = self._resource_manager.list_templates()
        return [
            MCPResourceTemplate(
                uriTemplate=template.uri_template,
                name=template.name,
                description=template.description,
            )
            for template in templates
        ]

    async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]:
        """Read a resource by URI."""

        resource = await self._resource_manager.get_resource(uri)
        if not resource:
            raise ResourceError(f"Unknown resource: {uri}")

        try:
            content = await resource.read()
            return [ReadResourceContents(content=content, mime_type=resource.mime_type)]
        except Exception as e:
            logger.error(f"Error reading resource {uri}: {e}")
            raise ResourceError(str(e))

    def add_tool(
        self,
        fn: AnyFunction,
        name: str | None = None,
        description: str | None = None,
        annotations: ToolAnnotations | None = None,
    ) -> None:
        """Add a tool to the server.

        The tool function can optionally request a Context object by adding a parameter
        with the Context type annotation. See the @tool decorator for examples.

        Args:
            fn: The function to register as a tool
            name: Optional name for the tool (defaults to function name)
            description: Optional description of what the tool does
            annotations: Optional ToolAnnotations providing additional tool information
        """
        self._tool_manager.add_tool(
            fn, name=name, description=description, annotations=annotations
        )

    def tool(
        self,
        name: str | None = None,
        description: str | None = None,
        annotations: ToolAnnotations | None = None,
        etdi: bool = False,
        etdi_permissions: list[str] | None = None,
        etdi_oauth_scopes: list[str] | None = None,
        etdi_max_call_depth: int | None = None,
        etdi_allowed_callees: list[str] | None = None,
        etdi_blocked_callees: list[str] | None = None,
        etdi_require_request_signing: bool = False,
    ) -> Callable[[AnyFunction], AnyFunction]:
        """Decorator to register a tool.

        Tools can optionally request a Context object by adding a parameter with the
        Context type annotation. The context provides access to MCP capabilities like
        logging, progress reporting, and resource access.

        Args:
            name: Optional name for the tool (defaults to function name)
            description: Optional description of what the tool does
            annotations: Optional ToolAnnotations providing additional tool information
            etdi: Enable ETDI (Enhanced Tool Definition Interface) security features
            etdi_permissions: List of permission scopes required for ETDI (e.g., ['data:read', 'files:write'])
            etdi_oauth_scopes: List of OAuth scopes for ETDI authentication
            etdi_max_call_depth: Maximum call stack depth for this tool
            etdi_allowed_callees: List of tool IDs this tool is allowed to call
            etdi_blocked_callees: List of tool IDs this tool is blocked from calling
            etdi_require_request_signing: Require cryptographic request signing (STRICT security level only)

        Example:
            @server.tool()
            def my_tool(x: int) -> str:
                return str(x)

            @server.tool()
            def tool_with_context(x: int, ctx: Context) -> str:
                ctx.info(f"Processing {x}")
                return str(x)

            @server.tool(etdi=True, etdi_permissions=['data:read'], etdi_max_call_depth=3)
            def secure_tool(x: int) -> str:
                return f"Securely processed: {x}"

            @server.tool(etdi=True, etdi_require_request_signing=True, etdi_permissions=['banking:write'])
            def ultra_secure_tool(amount: float) -> str:
                return f"Ultra-secure transaction: ${amount}"

            @server.tool()
            async def async_tool(x: int, context: Context) -> str:
                await context.report_progress(50, 100)
                return str(x)
        """
        # Check if user passed function directly instead of calling decorator
        if callable(name):
            raise TypeError(
                "The @tool decorator was used incorrectly. "
                "Did you forget to call it? Use @tool() instead of @tool"
            )

        def decorator(fn: AnyFunction) -> AnyFunction:
            # Handle ETDI integration
            if etdi and ETDI_AVAILABLE:
                # Lazy import ETDI components
                from mcp.etdi import ETDIToolDefinition, CallStackConstraints, Permission
                
                # Create ETDI tool definition
                tool_name = name or fn.__name__
                tool_description = description or fn.__doc__ or f"Tool: {tool_name}"
                
                # Create permissions from etdi_permissions
                permissions = []
                if etdi_permissions:
                    for perm_scope in etdi_permissions:
                        permissions.append(Permission(
                            name=perm_scope.replace(':', '_'),
                            description=f"Permission for {perm_scope}",
                            scope=perm_scope,
                            required=True
                        ))
                
                # Create call stack constraints if specified
                call_stack_constraints = None
                if any([etdi_max_call_depth, etdi_allowed_callees, etdi_blocked_callees]):
                    call_stack_constraints = CallStackConstraints(
                        max_depth=etdi_max_call_depth,
                        allowed_callees=etdi_allowed_callees,
                        blocked_callees=etdi_blocked_callees
                    )
                
                # Create ETDI tool definition
                etdi_tool = ETDIToolDefinition(
                    id=tool_name,
                    name=tool_name,
                    version="1.0.0",
                    description=tool_description,
                    provider={"id": "fastmcp", "name": "FastMCP Server"},
                    schema={"type": "object"},  # Will be filled by tool manager
                    permissions=permissions,
                    call_stack_constraints=call_stack_constraints
                )
                
                # Store ETDI metadata on the function for later use
                fn._etdi_tool_definition = etdi_tool
                fn._etdi_enabled = True
                fn._etdi_require_request_signing = etdi_require_request_signing
                
                # AUTOMATICALLY wrap the function with security enforcement
                fn = self._wrap_with_etdi_security(fn, etdi_tool, etdi_require_request_signing)
                
            elif etdi and not ETDI_AVAILABLE:
                # Warn if ETDI requested but not available
                import warnings
                warnings.warn(
                    f"ETDI requested for tool '{name or fn.__name__}' but ETDI is not available. "
                    "Install with 'pip install mcp[etdi]' to enable ETDI features.",
                    UserWarning
                )
                fn._etdi_enabled = False
            else:
                fn._etdi_enabled = False
            
            self.add_tool(
                fn, name=name, description=description, annotations=annotations
            )
            return fn

        return decorator
    
    def _wrap_with_etdi_security(self, fn: AnyFunction, etdi_tool: 'ETDIToolDefinition', require_request_signing: bool = False) -> AnyFunction:
        """Automatically wrap function with ETDI security enforcement"""
        if not ETDI_AVAILABLE:
            return fn
            
        def security_wrapper(*args, **kwargs):
            # 1. Check request signing (STRICT security level only)
            if require_request_signing:
                # Lazy import to avoid circular dependency
                from mcp.etdi.types import SecurityLevel
                from mcp.etdi.exceptions import SecurityError
                
                if self.settings.security_level == SecurityLevel.STRICT:
                    if not self._verify_request_signature():
                        raise SecurityError("Request signature verification failed. This tool requires cryptographic request signing.")
                elif self.settings.security_level != SecurityLevel.STRICT:
                    # Warn but don't block - backward compatibility
                    import warnings
                    warnings.warn(
                        f"Tool '{etdi_tool.name}' requires request signing but server is not in STRICT security mode. "
                        "Request signing is only enforced in STRICT mode for backward compatibility.",
                        UserWarning
                    )
            
            # 2. Check permissions
            if etdi_tool.permissions:
                required_perms = [p.scope for p in etdi_tool.permissions if p.required]
                if not self._check_permissions(required_perms):
                    missing = set(required_perms) - set(self._current_user_permissions)
                    raise PermissionError(f"Access denied. Missing permissions: {missing}")
            
            # 3. Verify call stack constraints
            if etdi_tool.call_stack_constraints and self._etdi_verifier:
                try:
                    self._etdi_verifier.verify_call(etdi_tool, session_id=self._current_session_id)
                except Exception as e:
                    raise SecurityError(f"ETDI security violation: {e}")
            
            # 4. Execute the original function if security checks pass
            return fn(*args, **kwargs)
        
        # Preserve function metadata
        security_wrapper.__name__ = fn.__name__
        security_wrapper.__doc__ = fn.__doc__
        security_wrapper._etdi_tool_definition = etdi_tool
        security_wrapper._etdi_enabled = True
        security_wrapper._etdi_require_request_signing = require_request_signing
        
        return security_wrapper
    
    def _check_permissions(self, required_permissions: list[str]) -> bool:
        """Check if current user has required permissions"""
        return all(perm in self._current_user_permissions for perm in required_permissions)
    
    def set_user_permissions(self, permissions: list[str]) -> None:
        """Set current user permissions (called by auth middleware)"""
        self._current_user_permissions = permissions
    
    def _verify_request_signature(self) -> bool:
        """Verify cryptographic signature of the current request"""
        if not ETDI_AVAILABLE or not hasattr(self, '_signature_verifier'):
            return False
        
        try:
            # Get current request context
            request_context = self.get_context().request_context
            if not request_context:
                return False
            
            # Get the current MCP request from context
            current_request = getattr(request_context, 'request', None)
            if not current_request:
                return False
            
            # Check if this is a CallToolRequest with ETDI signature headers
            if hasattr(current_request, 'params'):
                params = current_request.params
                
                # Extract ETDI signature headers from request parameters
                signature_headers = {}
                if hasattr(params, 'etdi_signature') and params.etdi_signature:
                    signature_headers['X-ETDI-Signature'] = params.etdi_signature
                if hasattr(params, 'etdi_timestamp') and params.etdi_timestamp:
                    signature_headers['X-ETDI-Timestamp'] = params.etdi_timestamp
                if hasattr(params, 'etdi_key_id') and params.etdi_key_id:
                    signature_headers['X-ETDI-Key-ID'] = params.etdi_key_id
                if hasattr(params, 'etdi_algorithm') and params.etdi_algorithm:
                    signature_headers['X-ETDI-Algorithm'] = params.etdi_algorithm
                
                # If no signature headers found, check if request has signature
                if not signature_headers:
                    return False
                
                # Verify the tool invocation signature
                is_valid = self._signature_verifier.verify_tool_invocation_signature(
                    tool_name=params.name,
                    arguments=params.arguments or {},
                    signature_headers=signature_headers
                )
                
                if not is_valid:
                    logger.warning(f"Tool invocation signature verification failed for {params.name}")
                else:
                    logger.debug(f"Tool invocation signature verified for {params.name}")
                
                return is_valid
            else:
                # No signature headers in request
                return False
            
        except Exception as e:
            logger.error(f"Error verifying request signature: {e}")
            return False
    
    def initialize_request_signing(self, key_store_path: str | None = None) -> None:
        """Initialize request signing verification (STRICT mode only)"""
        if not ETDI_AVAILABLE:
            return
        
        try:
            from mcp.etdi.crypto import KeyManager, SignatureVerifier
            
            key_manager = KeyManager(key_store_path)
            self._signature_verifier = SignatureVerifier(key_manager)
            self._key_manager = key_manager
            
            logger.info("Request signing verification initialized")
            
        except Exception as e:
            logger.error(f"Failed to initialize request signing: {e}")
            raise

    def add_resource(self, resource: Resource) -> None:
        """Add a resource to the server.

        Args:
            resource: A Resource instance to add
        """
        self._resource_manager.add_resource(resource)

    def resource(
        self,
        uri: str,
        *,
        name: str | None = None,
        description: str | None = None,
        mime_type: str | None = None,
    ) -> Callable[[AnyFunction], AnyFunction]:
        """Decorator to register a function as a resource.

        The function will be called when the resource is read to generate its content.
        The function can return:
        - str for text content
        - bytes for binary content
        - other types will be converted to JSON

        If the URI contains parameters (e.g. "resource://{param}") or the function
        has parameters, it will be registered as a template resource.

        Args:
            uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}")
            name: Optional name for the resource
            description: Optional description of the resource
            mime_type: Optional MIME type for the resource

        Example:
            @server.resource("resource://my-resource")
            def get_data() -> str:
                return "Hello, world!"

            @server.resource("resource://my-resource")
            async get_data() -> str:
                data = await fetch_data()
                return f"Hello, world! {data}"

            @server.resource("resource://{city}/weather")
            def get_weather(city: str) -> str:
                return f"Weather for {city}"

            @server.resource("resource://{city}/weather")
            async def get_weather(city: str) -> str:
                data = await fetch_weather(city)
                return f"Weather for {city}: {data}"
        """
        # Check if user passed function directly instead of calling decorator
        if callable(uri):
            raise TypeError(
                "The @resource decorator was used incorrectly. "
                "Did you forget to call it? Use @resource('uri') instead of @resource"
            )

        def decorator(fn: AnyFunction) -> AnyFunction:
            # Check if this should be a template
            has_uri_params = "{" in uri and "}" in uri
            has_func_params = bool(inspect.signature(fn).parameters)

            if has_uri_params or has_func_params:
                # Validate that URI params match function params
                uri_params = set(re.findall(r"{(\w+)}", uri))
                func_params = set(inspect.signature(fn).parameters.keys())

                if uri_params != func_params:
                    raise ValueError(
                        f"Mismatch between URI parameters {uri_params} "
                        f"and function parameters {func_params}"
                    )

                # Register as template
                self._resource_manager.add_template(
                    fn=fn,
                    uri_template=uri,
                    name=name,
                    description=description,
                    mime_type=mime_type,
                )
            else:
                # Register as regular resource
                resource = FunctionResource.from_function(
                    fn=fn,
                    uri=uri,
                    name=name,
                    description=description,
                    mime_type=mime_type,
                )
                self.add_resource(resource)
            return fn

        return decorator

    def add_prompt(self, prompt: Prompt) -> None:
        """Add a prompt to the server.

        Args:
            prompt: A Prompt instance to add
        """
        self._prompt_manager.add_prompt(prompt)

    def prompt(
        self, name: str | None = None, description: str | None = None
    ) -> Callable[[AnyFunction], AnyFunction]:
        """Decorator to register a prompt.

        Args:
            name: Optional name for the prompt (defaults to function name)
            description: Optional description of what the prompt does

        Example:
            @server.prompt()
            def analyze_table(table_name: str) -> list[Message]:
                schema = read_table_schema(table_name)
                return [
                    {
                        "role": "user",
                        "content": f"Analyze this schema:\n{schema}"
                    }
                ]

            @server.prompt()
            async def analyze_file(path: str) -> list[Message]:
                content = await read_file(path)
                return [
                    {
                        "role": "user",
                        "content": {
                            "type": "resource",
                            "resource": {
                                "uri": f"file://{path}",
                                "text": content
                            }
                        }
                    }
                ]
        """
        # Check if user passed function directly instead of calling decorator
        if callable(name):
            raise TypeError(
                "The @prompt decorator was used incorrectly. "
                "Did you forget to call it? Use @prompt() instead of @prompt"
            )

        def decorator(func: AnyFunction) -> AnyFunction:
            prompt = Prompt.from_function(func, name=name, description=description)
            self.add_prompt(prompt)
            return func

        return decorator

    def custom_route(
        self,
        path: str,
        methods: list[str],
        name: str | None = None,
        include_in_schema: bool = True,
    ):
        """
        Decorator to register a custom HTTP route on the FastMCP server.

        Allows adding arbitrary HTTP endpoints outside the standard MCP protocol,
        which can be useful for OAuth callbacks, health checks, or admin APIs.
        The handler function must be an async function that accepts a Starlette
        Request and returns a Response.

        Args:
            path: URL path for the route (e.g., "/oauth/callback")
            methods: List of HTTP methods to support (e.g., ["GET", "POST"])
            name: Optional name for the route (to reference this route with
                  Starlette's reverse URL lookup feature)
            include_in_schema: Whether to include in OpenAPI schema, defaults to True

        Example:
            @server.custom_route("/health", methods=["GET"])
            async def health_check(request: Request) -> Response:
                return JSONResponse({"status": "ok"})
        """

        def decorator(
            func: Callable[[Request], Awaitable[Response]],
        ) -> Callable[[Request], Awaitable[Response]]:
            self._custom_starlette_routes.append(
                Route(
                    path,
                    endpoint=func,
                    methods=methods,
                    name=name,
                    include_in_schema=include_in_schema,
                )
            )
            return func

        return decorator

    async def run_stdio_async(self) -> None:
        """Run the server using stdio transport."""
        async with stdio_server() as (read_stream, write_stream):
            await self._mcp_server.run(
                read_stream,
                write_stream,
                self._mcp_server.create_initialization_options(),
            )

    async def run_sse_async(self, mount_path: str | None = None) -> None:
        """Run the server using SSE transport."""
        import uvicorn

        starlette_app = self.sse_app(mount_path)

        config = uvicorn.Config(
            starlette_app,
            host=self.settings.host,
            port=self.settings.port,
            log_level=self.settings.log_level.lower(),
        )
        server = uvicorn.Server(config)
        await server.serve()

    async def run_streamable_http_async(self) -> None:
        """Run the server using StreamableHTTP transport."""
        import uvicorn

        starlette_app = self.streamable_http_app()

        config = uvicorn.Config(
            starlette_app,
            host=self.settings.host,
            port=self.settings.port,
            log_level=self.settings.log_level.lower(),
        )
        server = uvicorn.Server(config)
        await server.serve()

    def _normalize_path(self, mount_path: str, endpoint: str) -> str:
        """
        Combine mount path and endpoint to return a normalized path.

        Args:
            mount_path: The mount path (e.g. "/github" or "/")
            endpoint: The endpoint path (e.g. "/messages/")

        Returns:
            Normalized path (e.g. "/github/messages/")
        """
        # Special case: root path
        if mount_path == "/":
            return endpoint

        # Remove trailing slash from mount path
        if mount_path.endswith("/"):
            mount_path = mount_path[:-1]

        # Ensure endpoint starts with slash
        if not endpoint.startswith("/"):
            endpoint = "/" + endpoint

        # Combine paths
        return mount_path + endpoint

    def sse_app(self, mount_path: str | None = None) -> Starlette:
        """Return an instance of the SSE server app."""
        from starlette.middleware import Middleware
        from starlette.routing import Mount, Route

        # Update mount_path in settings if provided
        if mount_path is not None:
            self.settings.mount_path = mount_path

        # Create normalized endpoint considering the mount path
        normalized_message_endpoint = self._normalize_path(
            self.settings.mount_path, self.settings.message_path
        )

        # Set up auth context and dependencies

        sse = SseServerTransport(
            normalized_message_endpoint,
        )

        async def handle_sse(scope: Scope, receive: Receive, send: Send):
            # Add client ID from auth context into request context if available

            async with sse.connect_sse(
                scope,
                receive,
                send,
            ) as streams:
                await self._mcp_server.run(
                    streams[0],
                    streams[1],
                    self._mcp_server.create_initialization_options(),
                )
            return Response()

        # Create routes
        routes: list[Route | Mount] = []
        middleware: list[Middleware] = []
        required_scopes = []

        # Add auth endpoints if auth provider is configured
        if self._auth_server_provider:
            assert self.settings.auth
            from mcp.server.auth.routes import create_auth_routes

            required_scopes = self.settings.auth.required_scopes or []

            middleware = [
                # extract auth info from request (but do not require it)
                Middleware(
                    AuthenticationMiddleware,
                    backend=BearerAuthBackend(
                        provider=self._auth_server_provider,
                    ),
                ),
                # Add the auth context middleware to store
                # authenticated user in a contextvar
                Middleware(AuthContextMiddleware),
            ]
            routes.extend(
                create_auth_routes(
                    provider=self._auth_server_provider,
                    issuer_url=self.settings.auth.issuer_url,
                    service_documentation_url=self.settings.auth.service_documentation_url,
                    client_registration_options=self.settings.auth.client_registration_options,
                    revocation_options=self.settings.auth.revocation_options,
                )
            )

        # When auth is not configured, we shouldn't require auth
        if self._auth_server_provider:
            # Auth is enabled, wrap the endpoints with RequireAuthMiddleware
            routes.append(
                Route(
                    self.settings.sse_path,
                    endpoint=RequireAuthMiddleware(handle_sse, required_scopes),
                    methods=["GET"],
                )
            )
            routes.append(
                Mount(
                    self.settings.message_path,
                    app=RequireAuthMiddleware(sse.handle_post_message, required_scopes),
                )
            )
        else:
            # Auth is disabled, no need for RequireAuthMiddleware
            # Since handle_sse is an ASGI app, we need to create a compatible endpoint
            async def sse_endpoint(request: Request) -> Response:
                # Convert the Starlette request to ASGI parameters
                return await handle_sse(request.scope, request.receive, request._send)  # type: ignore[reportPrivateUsage]

            routes.append(
                Route(
                    self.settings.sse_path,
                    endpoint=sse_endpoint,
                    methods=["GET"],
                )
            )
            routes.append(
                Mount(
                    self.settings.message_path,
                    app=sse.handle_post_message,
                )
            )
        # mount these routes last, so they have the lowest route matching precedence
        routes.extend(self._custom_starlette_routes)

        # Create Starlette app with routes and middleware
        return Starlette(
            debug=self.settings.debug, routes=routes, middleware=middleware
        )

    def streamable_http_app(self) -> Starlette:
        """Return an instance of the StreamableHTTP server app."""
        from starlette.middleware import Middleware
        from starlette.routing import Mount

        # Create session manager on first call (lazy initialization)
        if self._session_manager is None:
            self._session_manager = StreamableHTTPSessionManager(
                app=self._mcp_server,
                event_store=self._event_store,
                json_response=self.settings.json_response,
                stateless=self.settings.stateless_http,  # Use the stateless setting
            )

        # Create the ASGI handler
        async def handle_streamable_http(
            scope: Scope, receive: Receive, send: Send
        ) -> None:
            await self.session_manager.handle_request(scope, receive, send)

        # Create routes
        routes: list[Route | Mount] = []
        middleware: list[Middleware] = []
        required_scopes = []

        # Add auth endpoints if auth provider is configured
        if self._auth_server_provider:
            assert self.settings.auth
            from mcp.server.auth.routes import create_auth_routes

            required_scopes = self.settings.auth.required_scopes or []

            middleware = [
                Middleware(
                    AuthenticationMiddleware,
                    backend=BearerAuthBackend(
                        provider=self._auth_server_provider,
                    ),
                ),
                Middleware(AuthContextMiddleware),
            ]
            routes.extend(
                create_auth_routes(
                    provider=self._auth_server_provider,
                    issuer_url=self.settings.auth.issuer_url,
                    service_documentation_url=self.settings.auth.service_documentation_url,
                    client_registration_options=self.settings.auth.client_registration_options,
                    revocation_options=self.settings.auth.revocation_options,
                )
            )
            routes.append(
                Mount(
                    self.settings.streamable_http_path,
                    app=RequireAuthMiddleware(handle_streamable_http, required_scopes),
                )
            )
        else:
            # Auth is disabled, no wrapper needed
            routes.append(
                Mount(
                    self.settings.streamable_http_path,
                    app=handle_streamable_http,
                )
            )

        routes.extend(self._custom_starlette_routes)

        return Starlette(
            debug=self.settings.debug,
            routes=routes,
            middleware=middleware,
            lifespan=lambda app: self.session_manager.run(),
        )

    async def list_prompts(self) -> list[MCPPrompt]:
        """List all available prompts."""
        prompts = self._prompt_manager.list_prompts()
        return [
            MCPPrompt(
                name=prompt.name,
                description=prompt.description,
                arguments=[
                    MCPPromptArgument(
                        name=arg.name,
                        description=arg.description,
                        required=arg.required,
                    )
                    for arg in (prompt.arguments or [])
                ],
            )
            for prompt in prompts
        ]

    async def get_prompt(
        self, name: str, arguments: dict[str, Any] | None = None
    ) -> GetPromptResult:
        """Get a prompt by name with arguments."""
        try:
            messages = await self._prompt_manager.render_prompt(name, arguments)

            return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages))
        except Exception as e:
            logger.error(f"Error getting prompt {name}: {e}")
            raise ValueError(str(e))


def _convert_to_content(
    result: Any,
) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
    """Convert a result to a sequence of content objects."""
    if result is None:
        return []

    if isinstance(result, TextContent | ImageContent | EmbeddedResource):
        return [result]

    if isinstance(result, Image):
        return [result.to_image_content()]

    if isinstance(result, list | tuple):
        return list(chain.from_iterable(_convert_to_content(item) for item in result))  # type: ignore[reportUnknownVariableType]

    if not isinstance(result, str):
        result = pydantic_core.to_json(result, fallback=str, indent=2).decode()

    return [TextContent(type="text", text=result)]


class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
    """Context object providing access to MCP capabilities.

    This provides a cleaner interface to MCP's RequestContext functionality.
    It gets injected into tool and resource functions that request it via type hints.

    To use context in a tool function, add a parameter with the Context type annotation:

    ```python
    @server.tool()
    def my_tool(x: int, ctx: Context) -> str:
        # Log messages to the client
        ctx.info(f"Processing {x}")
        ctx.debug("Debug info")
        ctx.warning("Warning message")
        ctx.error("Error message")

        # Report progress
        ctx.report_progress(50, 100)

        # Access resources
        data = ctx.read_resource("resource://data")

        # Get request info
        request_id = ctx.request_id
        client_id = ctx.client_id

        return str(x)
    ```

    The context parameter name can be anything as long as it's annotated with Context.
    The context is optional - tools that don't need it can omit the parameter.
    """

    _request_context: RequestContext[ServerSessionT, LifespanContextT] | None
    _fastmcp: FastMCP | None

    def __init__(
        self,
        *,
        request_context: RequestContext[ServerSessionT, LifespanContextT] | None = None,
        fastmcp: FastMCP | None = None,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self._request_context = request_context
        self._fastmcp = fastmcp

    @property
    def fastmcp(self) -> FastMCP:
        """Access to the FastMCP server."""
        if self._fastmcp is None:
            raise ValueError("Context is not available outside of a request")
        return self._fastmcp

    @property
    def request_context(self) -> RequestContext[ServerSessionT, LifespanContextT]:
        """Access to the underlying request context."""
        if self._request_context is None:
            raise ValueError("Context is not available outside of a request")
        return self._request_context

    async def report_progress(
        self, progress: float, total: float | None = None, message: str | None = None
    ) -> None:
        """Report progress for the current operation.

        Args:
            progress: Current progress value e.g. 24
            total: Optional total value e.g. 100
            message: Optional message e.g. Starting render...
        """
        progress_token = (
            self.request_context.meta.progressToken
            if self.request_context.meta
            else None
        )

        if progress_token is None:
            return

        await self.request_context.session.send_progress_notification(
            progress_token=progress_token,
            progress=progress,
            total=total,
            message=message,
        )

    async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]:
        """Read a resource by URI.

        Args:
            uri: Resource URI to read

        Returns:
            The resource content as either text or bytes
        """
        assert (
            self._fastmcp is not None
        ), "Context is not available outside of a request"
        return await self._fastmcp.read_resource(uri)

    async def log(
        self,
        level: Literal["debug", "info", "warning", "error"],
        message: str,
        *,
        logger_name: str | None = None,
    ) -> None:
        """Send a log message to the client.

        Args:
            level: Log level (debug, info, warning, error)
            message: Log message
            logger_name: Optional logger name
            **extra: Additional structured data to include
        """
        await self.request_context.session.send_log_message(
            level=level,
            data=message,
            logger=logger_name,
            related_request_id=self.request_id,
        )

    @property
    def client_id(self) -> str | None:
        """Get the client ID if available."""
        return (
            getattr(self.request_context.meta, "client_id", None)
            if self.request_context.meta
            else None
        )

    @property
    def request_id(self) -> str:
        """Get the unique ID for this request."""
        return str(self.request_context.request_id)

    @property
    def session(self):
        """Access to the underlying session for advanced usage."""
        return self.request_context.session

    # Convenience methods for common log levels
    async def debug(self, message: str, **extra: Any) -> None:
        """Send a debug log message."""
        await self.log("debug", message, **extra)

    async def info(self, message: str, **extra: Any) -> None:
        """Send an info log message."""
        await self.log("info", message, **extra)

    async def warning(self, message: str, **extra: Any) -> None:
        """Send a warning log message."""
        await self.log("warning", message, **extra)

    async def error(self, message: str, **extra: Any) -> None:
        """Send an error log message."""
        await self.log("error", message, **extra)
