"""Internal MCP aggregator — one ASGI app, N FastMCP namespaces + ``all``.

Each ``McpServer`` declared in the user's config becomes its own
``FastMCP`` instance (regular for ``python_tool``, ``FastMCPProxy`` for
``stdio``/``http``) and is mounted under ``/mcp/<name>`` on a single
Starlette app. This app runs on ``127.0.0.1:INTERNAL_MCP_PORT`` (not
EXPOSE'd in Docker) so the ClaudeCode subprocess can reach each
namespace via loopback as a distinct MCP server URL — preserving
per-domain framing while costing only one process worth of RAM
(PRD §6).

Phase 3 adds ``/mcp/all/``: a single FastMCP whose tools are the union
of every namespace's tools, prefixed by FastMCP's ``namespace_<tool>``
convention (e.g. ``time_current_time``). It's the escape-hatch for
clients that can only configure one MCP server — discouraged for tool-
heavy setups (PRD §6 cites the ~95%→~71% tool-selection drop on flat
namespaces) but real and reachable.

The aggregator returns both the app and a ``{name: url}`` map; Phase
``ClaudeSdkBackend`` plugs the map directly into
``BackendOptions.mcp_servers``. ``/mcp/all/`` is NOT included in that
map — claude-code-agents always get per-domain framing; only the
external MCP frontend (Phase 3.1) reverse-proxies the flat endpoint.
"""

from __future__ import annotations

from contextlib import AsyncExitStack, asynccontextmanager
from typing import TYPE_CHECKING

from fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.routing import Mount

from beaver_gateway.mcp.client_pool import build_http_proxy, build_stdio_proxy
from beaver_gateway.mcp.types import HttpMcp, PythonToolMcp, StdioMcp
from beaver_gateway.mcp.wrap import build_python_tool_server

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Iterable

    from beaver_gateway.mcp.types import McpServerT


ALL_NAMESPACE = "all"
"""URL segment for the flat-namespace aggregator (``/mcp/all/``)."""


def build_internal_app(
    mcps: Iterable[McpServerT], *, host: str, port: int
) -> tuple[Starlette, dict[str, str], dict[str, FastMCP]]:
    """Build the aggregator ``Starlette`` app, per-namespace URL map, and server map.

    ``host``/``port`` only flavour the URL strings handed back — actually
    listening on them is the caller's job (``cli.main`` runs a uvicorn
    server in a TaskGroup). We accept the address here so callers don't
    have to format the URLs themselves and risk drifting from the
    ``/mcp/<name>`` convention.

    Returns:
      * Starlette app to serve via uvicorn.
      * ``{namespace: url}`` for the per-domain endpoints (claude-code's
        MCP routing expects per-domain framing). ``/mcp/all/`` is
        omitted — only meaningful to external clients via the MCP
        frontend, not to claude-code.
      * ``{namespace: FastMCP}`` for backends that need in-process
        access (Raycast doesn't natively understand MCP, so the
        gateway calls ``list_tools``/``call_tool`` directly to splice
        MCP tools into the Raycast wire).
    """
    servers: dict[str, FastMCP] = {spec.name: _build_server(spec) for spec in mcps}

    child_apps = {
        name: s.http_app(transport="http", path="/") for name, s in servers.items()
    }
    routes = [Mount(f"/mcp/{name}", app=app) for name, app in child_apps.items()]

    # /mcp/all — flat-namespace bundle. Skip when there's nothing to
    # bundle so we don't pay for an empty session manager lifecycle.
    all_app = None
    if servers:
        all_server = _build_all_server(servers)
        all_app = all_server.http_app(transport="http", path="/")
        routes.append(Mount(f"/mcp/{ALL_NAMESPACE}", app=all_app))

    @asynccontextmanager
    async def lifespan(_parent: Starlette) -> AsyncIterator[None]:
        # Each FastMCP http_app stores its session manager init in its
        # own lifespan. Without entering them the streamable-HTTP layer
        # 500s on every request. AsyncExitStack composes them so all
        # children come up together and unwind in reverse order on
        # shutdown.
        async with AsyncExitStack() as stack:
            for child in child_apps.values():
                await stack.enter_async_context(child.router.lifespan_context(child))
            if all_app is not None:
                await stack.enter_async_context(
                    all_app.router.lifespan_context(all_app)
                )
            yield

    app = Starlette(routes=routes, lifespan=lifespan)
    # Trailing slash on the published URL skips Starlette's
    # 307 redirect from ``/mcp/<name>`` to ``/mcp/<name>/`` that
    # ``Mount`` produces when a child route lives at ``/``.
    urls = {name: f"http://{host}:{port}/mcp/{name}/" for name in servers}
    return app, urls, servers


def _build_server(spec: McpServerT) -> FastMCP:
    """Dispatch on the discriminated union to the matching builder."""
    if isinstance(spec, PythonToolMcp):
        return build_python_tool_server(spec)
    if isinstance(spec, StdioMcp):
        return build_stdio_proxy(spec)
    if isinstance(spec, HttpMcp):
        return build_http_proxy(spec)
    # `McpServerT` is a closed union; this is unreachable but keeps
    # type-narrowing honest if a new variant lands without updates here.
    msg = f"unsupported McpServer variant: {type(spec).__name__}"
    raise TypeError(msg)


def _build_all_server(children: dict[str, FastMCP]) -> FastMCP:
    """Compose a single FastMCP whose tools union every namespace's tools.

    FastMCP's ``mount(namespace=...)`` namespaces every tool as
    ``<namespace>_<original_name>``, so the flat endpoint becomes
    ``time_current_time``, ``calendar_event_create``, etc.

    Mounted children are accessed in-memory via ``FastMCPProvider`` —
    requests don't bounce through the child's own ``http_app``, so this
    aggregator has its own independent streamable-HTTP session manager
    and lifespan, and the per-namespace ``/mcp/<name>/`` mounts keep
    working unchanged.
    """
    parent = FastMCP(name="beaver-gateway-all")
    for name, child in children.items():
        parent.mount(child, namespace=name)
    return parent