# Copyright (c) 2026 Dedalus Labs, Inc. and its contributors
# SPDX-License-Identifier: MIT

"""ASGI transport primitives for MCP servers."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Callable, Iterable, Mapping  # noqa: TC003
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol

from starlette.applications import Starlette
from uvicorn import Config, Server

from .base import BaseTransport


if TYPE_CHECKING:
    from starlette.routing import BaseRoute
    from starlette.types import Receive, Scope, Send

    from ..authorization import AuthorizationManager
    from ..core import MCPServer


@dataclass(slots=True)
class ASGITransportConfig:
    """Configuration toggles that shape transport behaviour."""

    security_settings: object | None = None
    stateless: bool = False


@dataclass(slots=True)
class ASGIRunConfig:
    """Runtime parameters for launching an ASGI transport."""

    host: str | None = None
    port: int | None = None
    path: str | None = None
    log_level: str | None = None
    uvicorn_options: Mapping[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class _ResolvedRunConfig:
    """Internal helper that holds concrete ASGI run settings."""

    host: str
    port: int
    path: str
    log_level: str
    uvicorn_options: dict[str, Any]


class SessionManagerProtocol(Protocol):
    """Minimal contract required of the reference SDK session managers."""

    async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None: ...

    def run(self) -> AbstractAsyncContextManager[None]: ...


@dataclass(slots=True)
class SessionManagerHandler:
    """ASGI adapter that connects the server session manager to the runtime."""

    session_manager: SessionManagerProtocol
    transport_label: str
    allowed_scopes: tuple[str, ...]

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        scope_type = scope.get("type")
        if scope_type not in self.allowed_scopes:
            allowed = ", ".join(self.allowed_scopes)
            message = f"{self.transport_label} only handles ASGI scopes: {allowed} (got {scope_type!r})."
            raise TypeError(message)

        await self.session_manager.handle_request(scope, receive, send)

    def lifespan(self) -> Callable[[Starlette], AbstractAsyncContextManager[None]]:
        """Return an ASGI lifespan hook bound to the session manager."""

        @asynccontextmanager
        async def _lifespan(_app: Starlette) -> AsyncIterator[None]:
            async with self.session_manager.run():
                yield

        return _lifespan


class ASGITransportBase(BaseTransport, ABC):
    """Template for transports that present an :class:`MCPServer` via ASGI."""

    ALLOWED_SCOPES: tuple[str, ...] = ("http",)
    DEFAULT_HOST: str = "127.0.0.1"
    DEFAULT_PORT: int = 8000
    DEFAULT_PATH: str = "/mcp"
    DEFAULT_LOG_LEVEL: str = "info"

    def __init__(self, server: MCPServer, *, config: ASGITransportConfig | None = None) -> None:
        super().__init__(server)
        self._config = config or ASGITransportConfig()
        self._server_instance: Server | None = None

    @property
    def security_settings(self) -> object | None:
        """Return the transport-specific security configuration, if any."""
        return self._config.security_settings

    @property
    def stateless(self) -> bool:
        """Return ``True`` when incoming requests should be treated statelessly."""
        return self._config.stateless

    async def run(self, *, config: ASGIRunConfig | None = None, **legacy_kwargs: Any) -> None:
        resolved = self._resolve_run_config(config=config, legacy_kwargs=legacy_kwargs)
        await self._serve(resolved)

    def _resolve_run_config(self, *, config: ASGIRunConfig | None, legacy_kwargs: dict[str, Any]) -> _ResolvedRunConfig:
        if config is not None and legacy_kwargs:
            raise TypeError("Cannot mix 'config' with legacy keyword arguments.")

        if config is None:
            if legacy_kwargs:
                extracted = {
                    k: legacy_kwargs.pop(k) for k in ("host", "port", "path", "log_level") if k in legacy_kwargs
                }
                config = ASGIRunConfig(
                    host=extracted.get("host"),
                    port=extracted.get("port"),
                    path=extracted.get("path"),
                    log_level=extracted.get("log_level"),
                    uvicorn_options=dict(legacy_kwargs),
                )
            else:
                config = ASGIRunConfig()
        else:
            # Normalize provided uvicorn options into a concrete dict we can mutate safely.
            legacy_kwargs = {}

        host = config.host or self.DEFAULT_HOST
        port = config.port or self.DEFAULT_PORT
        path = config.path or self.DEFAULT_PATH
        log_level = config.log_level or self.DEFAULT_LOG_LEVEL

        extra_options = dict(config.uvicorn_options)
        extra_options.update(legacy_kwargs)

        return _ResolvedRunConfig(host=host, port=port, path=path, log_level=log_level, uvicorn_options=extra_options)

    async def _serve(self, run_config: _ResolvedRunConfig) -> None:
        manager = self._build_session_manager()
        handler = self._build_handler(manager)
        routes = list(self._build_routes(path=run_config.path, handler=handler))

        authorization: AuthorizationManager | None = getattr(self.server, "authorization_manager", None)
        if authorization and authorization.enabled:
            routes.append(authorization.starlette_route())

        lifespan = handler.lifespan()
        asgi_app = Starlette(routes=routes, lifespan=lifespan)

        app = self._to_asgi(asgi_app)
        if authorization and authorization.enabled:
            app = authorization.wrap_asgi(app)

        uvicorn_config = Config(
            app=app,
            host=run_config.host,
            port=run_config.port,
            log_level=run_config.log_level,
            **run_config.uvicorn_options,
        )
        server_instance = Server(uvicorn_config)
        self._server_instance = server_instance
        try:
            await server_instance.serve()
        finally:
            self._server_instance = None

    async def stop(self) -> None:
        server_instance = self._server_instance
        if server_instance is None:
            return
        server_instance.should_exit = True
        await server_instance.shutdown()

    def _build_handler(self, manager: SessionManagerProtocol) -> SessionManagerHandler:
        """Construct the default ASGI handler for the provided session manager."""
        return SessionManagerHandler(
            session_manager=manager, transport_label=self.transport_display_name, allowed_scopes=self.ALLOWED_SCOPES
        )

    def _to_asgi(self, app: Starlette) -> Starlette:
        """Allow subclasses to wrap the ASGI app before serving.

        For example, a transport could override this method to inject
        instrumentation middleware before handing control to the ASGI server
        runtime.
        """
        return app

    @abstractmethod
    def _build_session_manager(self) -> SessionManagerProtocol: ...

    @abstractmethod
    def _build_routes(self, *, path: str, handler: SessionManagerHandler) -> Iterable[BaseRoute]: ...


__all__ = ["ASGITransportBase", "ASGITransportConfig", "ASGIRunConfig", "SessionManagerHandler"]
