"""External MCP frontend (Phase 3.1).

A streamable-HTTP gateway in front of the internal MCP aggregator
(``beaver_gateway.mcp.internal_app``). The aggregator hosts every
declared ``McpServer`` (``python_tool``, stdio proxy, HTTP proxy)
under ``/mcp/<name>`` plus a flat ``/mcp/all`` bundle on
``127.0.0.1:INTERNAL_MCP_PORT`` — that's the *internal* shape.

This frontend re-exposes those namespaces on its own port directly at
``/<name>/`` (no ``/mcp/`` prefix in the external routes — the port
itself already disambiguates). Caddy / nginx / Cloudflare in front
typically strips a prefix back on: ``domain.com/mcp/* → :8001/*``,
controlled by the operator's reverse-proxy config and surfaced to the
admin dashboard via ``public_base_url``. Three additions on top of the
raw aggregator:

* **Bearer auth** — ``Authorization: Bearer <token>``, ``X-Api-Key``,
  or ``?token=<…>`` query string. All three forms verify against the
  same :class:`TokenStore` as ``AnthropicMessagesFrontend``.
* **Audit log** — one line per request (token name, namespace,
  request method/path, response status). The DB-backed audit log lives
  in Phase 4; for now we just emit a structured log line.
* **Discovery page** at ``GET /`` (auth-gated) — HTML rendered with a
  tiny inline Jinja2 template listing every namespace plus copy-pastable
  config snippets for Cursor / claude.ai / Claude Desktop.

Why a reverse proxy and not a second mount? FastMCP's session managers
are tied to the lifespan they were created in; running the same
aggregator under two uvicorn servers double-initializes state. Building
two parallel aggregators would double upstream connections (two
subprocesses for every stdio MCP, two HTTP clients for every remote).
A loopback proxy keeps one source of truth — the internal aggregator —
and lets us layer policy on the outside.
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING, Any
from urllib.parse import urlencode

import aiohttp
from starlette.applications import Starlette
from starlette.responses import HTMLResponse, JSONResponse, StreamingResponse
from starlette.routing import Route

from beaver_gateway.core import audit
from beaver_gateway.frontends._urls import external_base
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.mcp.internal_app import ALL_NAMESPACE

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Mapping

    from starlette.requests import Request

    from beaver_gateway.frontends.base import GatewayRuntime


_log = logging.getLogger("beaver_gateway.frontends.mcp_server")


# Hop-by-hop headers that must NOT be forwarded across an HTTP proxy
# (RFC 7230 §6.1). Bypassing this filter would break chunked transfer
# encoding when ``Content-Length`` arrives, or upstream-aware proxies
# would refuse the second hop's connection-pool reuse.
_HOP_BY_HOP_HEADERS = frozenset(
    {
        "connection",
        "keep-alive",
        "proxy-authenticate",
        "proxy-authorization",
        "te",
        "trailers",
        "transfer-encoding",
        "upgrade",
        "host",
        "content-length",
    }
)

# Standard auth-bearing headers we *do not* forward to the internal app —
# the internal app is on loopback with no auth of its own, and forwarding
# the inbound bearer would only confuse it. Each method-specific MCP
# request from the upstream Cursor/etc. carries a fresh ``mcp-session-id``
# that we *must* forward.
_AUTH_HEADERS = frozenset({"authorization", "x-api-key"})


__all__ = ["McpServerFrontend"]


class McpServerFrontend(Frontend):
    """Auth + audit + reverse-proxy in front of the internal MCP aggregator."""

    path = "/mcp"

    def __init__(self) -> None:
        self._runtime: GatewayRuntime | None = None
        self._app: Starlette | None = None
        self._http: aiohttp.ClientSession | None = None

    def configure(self, runtime: GatewayRuntime) -> None:
        self._runtime = runtime
        self._app = self._build_app(runtime)

    def app(self) -> Starlette | None:
        return self._app

    async def serve(self) -> None:
        self._http = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=None, sock_read=600)
        )
        try:
            await asyncio.Event().wait()
        finally:
            await self._http.close()
            self._http = None

    def _build_app(self, runtime: GatewayRuntime) -> Starlette:  # noqa: ARG002
        routes = [
            Route("/", self._discovery, methods=["GET"]),
            Route("/healthz", self._healthz, methods=["GET"]),
            # Namespaces mount at the root of this port — the port
            # itself already disambiguates this from any other gateway
            # surface. Two routes per namespace so both the
            # trailing-slash and sub-path forms work (``/time`` AND
            # ``/time/foo``); Starlette doesn't fold them into one
            # route automatically. The literal routes above (``/``,
            # ``/healthz``) are listed first and win the match, so
            # they're not eaten by ``/{namespace}``.
            Route(
                "/{namespace}",
                self._proxy_endpoint,
                methods=["GET", "POST", "DELETE", "OPTIONS"],
            ),
            Route(
                "/{namespace}/{path:path}",
                self._proxy_endpoint,
                methods=["GET", "POST", "DELETE", "OPTIONS"],
            ),
        ]
        return Starlette(routes=routes)

    async def _healthz(self, _request: Request) -> JSONResponse:
        return JSONResponse({"status": "ok"})

    async def _discovery(self, request: Request) -> HTMLResponse | JSONResponse:
        runtime = self._require_runtime()
        token_name, err = await _verify_request(request, runtime)
        if err is not None:
            return err
        assert token_name is not None  # noqa: S101 — narrow for ty
        base = external_base(request, runtime)
        html = _render_discovery_page(
            base_url=base, namespaces=list(runtime.mcps), actor=token_name
        )
        return HTMLResponse(html)

    async def _proxy_endpoint(
        self, request: Request
    ) -> StreamingResponse | JSONResponse:
        runtime = self._require_runtime()
        token_name, err = await _verify_request(request, runtime)
        if err is not None:
            return err
        assert token_name is not None  # noqa: S101 — narrow for ty

        namespace = request.path_params["namespace"]
        subpath = request.path_params.get("path", "")
        upstream_url = self._upstream_url(namespace, subpath)
        if upstream_url is None:
            _log.info(
                "mcp: actor=%s namespace=%s 404 (unknown namespace)",
                token_name,
                namespace,
            )
            await audit.log(
                runtime,
                actor=f"token:{token_name}",
                kind="mcp_call",
                namespace=namespace,
                method=request.method,
                status=404,
            )
            return JSONResponse(
                {"error": "unknown namespace", "namespace": namespace}, status_code=404
            )

        if self._http is None:
            return JSONResponse({"error": "frontend not ready"}, status_code=503)

        return await _reverse_proxy(
            client=self._http,
            request=request,
            upstream_url=upstream_url,
            namespace=namespace,
            actor=token_name,
            runtime=runtime,
        )

    def _upstream_url(self, namespace: str, subpath: str) -> str | None:
        runtime = self._require_runtime()
        # ``all`` is built by the aggregator unconditionally when at least
        # one MCP is configured; the URL map only contains per-domain
        # entries (see ``build_internal_app``), so we synthesize ``all``'s
        # loopback URL from any per-domain URL's authority.
        if namespace == ALL_NAMESPACE:
            sample = next(iter(runtime.mcp_internal_urls.values()), None)
            if sample is None:
                return None
            base = sample.rsplit("/mcp/", 1)[0]
            return _join_subpath(f"{base}/mcp/{ALL_NAMESPACE}/", subpath)
        base = runtime.mcp_internal_urls.get(namespace)
        if base is None:
            return None
        return _join_subpath(base, subpath)

    def _require_runtime(self) -> GatewayRuntime:
        if self._runtime is None:
            msg = "configure() must be called before serving requests"
            raise RuntimeError(msg)
        return self._runtime


_MCP_SCOPE = "mcp"


async def _verify_request(
    request: Request, runtime: GatewayRuntime
) -> tuple[str | None, JSONResponse | None]:
    """Accept ``Authorization: Bearer``, ``X-Api-Key``, or ``?token=``.

    The third form is the escape hatch for clients that can only put
    secrets in the URL (claude.ai's MCP config historically did this).
    All three roads end at the same :class:`TokenStore`. Returns
    ``(actor_name, None)`` on success, ``(None, 401|403)`` otherwise
    — the caller forwards the response as-is. Splitting auth vs scope
    failures matters: 401 says "send me a token", 403 says "this token
    is real but not for this endpoint".
    """
    api_key = request.headers.get("x-api-key")
    if api_key:
        identity = await runtime.token_store.verify(api_key)
    else:
        auth_header = request.headers.get("authorization")
        if auth_header:
            identity = await runtime.token_store.verify_bearer(auth_header)
        else:
            qs_token = request.query_params.get("token")
            identity = await runtime.token_store.verify(qs_token) if qs_token else None

    if identity is None:
        return None, _unauthorized()
    if not identity.allows(_MCP_SCOPE):
        return None, _forbidden(identity.scope, _MCP_SCOPE)
    return identity.name, None


def _unauthorized() -> JSONResponse:
    return JSONResponse(
        {"error": "invalid or missing bearer token"},
        status_code=401,
        headers={"WWW-Authenticate": "Bearer"},
    )


def _forbidden(scope: str, required: str) -> JSONResponse:
    return JSONResponse(
        {"error": "insufficient scope", "scope": scope, "required": required},
        status_code=403,
    )


def _join_subpath(base_url: str, subpath: str) -> str:
    """Concatenate the loopback URL with the proxied sub-path.

    ``base_url`` always ends in ``/`` (the aggregator publishes URLs
    that way to avoid Starlette's 307 redirect dance); the sub-path is
    appended verbatim, with the query string handled by the caller.
    """
    if subpath:
        return base_url + subpath.lstrip("/")
    return base_url


async def _reverse_proxy(
    *,
    client: aiohttp.ClientSession,
    request: Request,
    upstream_url: str,
    namespace: str,
    actor: str,
    runtime: GatewayRuntime,
) -> StreamingResponse | JSONResponse:
    """Bidirectionally stream an MCP request between client ↔ internal aggregator.

    Streamable-HTTP MCP responses can be a long-running SSE stream
    (tools that emit partial progress) or a one-shot JSON body; we
    don't peek — just relay chunks as they arrive in either direction
    until both sides close.
    """
    qs = request.url.query
    if qs:
        # Drop ``?token=`` from the forwarded URL — internal app doesn't
        # need it, and propagating creds further than necessary widens
        # the leak surface (logs, metrics, traces all see query strings).
        scrubbed = _scrub_query(qs, drop={"token"})
        if scrubbed:
            upstream_url = f"{upstream_url}?{scrubbed}"

    forward_headers = _forward_headers(request)

    body_iter: AsyncIterator[bytes] | None = None
    if request.method not in {"GET", "HEAD", "OPTIONS"}:
        body_iter = _request_body_iter(request)

    try:
        upstream_resp = await client.request(
            request.method,
            upstream_url,
            headers=forward_headers,
            data=body_iter,
            allow_redirects=False,
        )
    except aiohttp.ClientError as exc:
        _log.warning(
            "mcp: actor=%s namespace=%s upstream connect failed: %s",
            actor,
            namespace,
            exc,
        )
        await audit.log(
            runtime,
            actor=f"token:{actor}",
            kind="mcp_call",
            namespace=namespace,
            method=request.method,
            status=502,
        )
        return JSONResponse(
            {"error": "upstream MCP unreachable", "detail": str(exc)}, status_code=502
        )

    _log.info(
        "mcp: actor=%s namespace=%s %s %s -> %d",
        actor,
        namespace,
        request.method,
        request.url.path,
        upstream_resp.status,
    )
    # Audit at upstream-response time: status reflects the MCP call's
    # outcome (200 / tool-error / 4xx). Streaming relay below may be
    # cut short by the client, but the row is already in by then.
    await audit.log(
        runtime,
        actor=f"token:{actor}",
        kind="mcp_call",
        namespace=namespace,
        method=request.method,
        status=upstream_resp.status,
    )

    response_headers = _response_headers(upstream_resp.headers)

    async def relay() -> AsyncIterator[bytes]:
        try:
            async for chunk in upstream_resp.content.iter_any():
                yield chunk
        except (aiohttp.ClientError, asyncio.CancelledError):
            # Caller hung up or upstream dropped — just stop relaying.
            return
        finally:
            upstream_resp.release()

    return StreamingResponse(
        relay(),
        status_code=upstream_resp.status,
        headers=response_headers,
        media_type=upstream_resp.headers.get("content-type"),
    )


async def _request_body_iter(request: Request) -> AsyncIterator[bytes]:
    async for chunk in request.stream():
        if chunk:
            yield chunk


def _forward_headers(request: Request) -> dict[str, str]:
    out: dict[str, str] = {}
    for key, value in request.headers.items():
        lowered = key.lower()
        if lowered in _HOP_BY_HOP_HEADERS or lowered in _AUTH_HEADERS:
            continue
        out[key] = value
    return out


def _response_headers(headers: Mapping[str, str]) -> dict[str, str]:
    out: dict[str, str] = {}
    for key, value in headers.items():
        if key.lower() in _HOP_BY_HOP_HEADERS:
            continue
        out[key] = value
    return out


def _scrub_query(query: str, *, drop: frozenset[str] | set[str]) -> str:
    """Remove sensitive keys from a URL-encoded query string."""
    kept: list[tuple[str, str]] = []
    for entry in query.split("&"):
        if not entry:
            continue
        name, sep, value = entry.partition("=")
        if name in drop:
            continue
        kept.append((name, value if sep else ""))
    return urlencode(kept)


def _render_discovery_page(*, base_url: str, namespaces: list[Any], actor: str) -> str:
    """Render the auth-gated namespace + config-snippet page.

    Inline HTML (no Jinja file) — keeps Phase 3 free of template-dir
    plumbing that Phase 4's AdminFrontend will own.
    """
    name_list = [getattr(ns, "name", str(ns)) for ns in namespaces]
    rows = (
        "\n".join(
            f"""        <tr>
          <td><code>{_escape(name)}</code></td>
          <td><code>{_escape(f"{base_url}/{name}/")}</code></td>
        </tr>"""
            for name in name_list
        )
        or """        \
<tr><td colspan="2"><em>No MCP servers configured.</em></td></tr>"""
    )
    cursor_snippet = _CURSOR_SNIPPET.format(base_url=base_url)
    claude_desktop_snippet = _CLAUDE_DESKTOP_SNIPPET.format(base_url=base_url)
    return _DISCOVERY_TEMPLATE.format(
        actor=_escape(actor),
        base_url=_escape(base_url),
        rows=rows,
        all_url=_escape(f"{base_url}/{ALL_NAMESPACE}/"),
        cursor_snippet=_escape(cursor_snippet),
        claude_desktop_snippet=_escape(claude_desktop_snippet),
    )


def _escape(value: str) -> str:
    return (
        value.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


_DISCOVERY_TEMPLATE = "\n".join(  # noqa: FLY002 — readability beats one-string-blob
    [
        "<!doctype html>",
        '<html lang="en">',
        "<head>",
        '  <meta charset="utf-8">',
        "  <title>beaver-gateway · MCP discovery</title>",
        "  <style>",
        "    body {{",
        "      font-family: -apple-system, 'SF Pro Display', system-ui, sans-serif;",
        "      max-width: 880px;",
        "      margin: 3rem auto;",
        "      padding: 0 1.25rem;",
        "      color: #1d1d1f;",
        "      line-height: 1.55;",
        "    }}",
        "    h1 {{ font-weight: 600; letter-spacing: -0.02em; }}",
        "    code, pre {{",
        "      font-family: ui-monospace, 'SF Mono', Menlo, monospace;",
        "      font-size: 0.92em;",
        "    }}",
        "    pre {{",
        "      background: #f5f5f7;",
        "      padding: 1rem 1.25rem;",
        "      border-radius: 10px;",
        "      overflow-x: auto;",
        "    }}",
        "    table {{ width: 100%; border-collapse: collapse; }}",
        "    th, td {{",
        "      text-align: left;",
        "      padding: 0.55rem 0.85rem;",
        "      border-bottom: 1px solid #e5e5ea;",
        "    }}",
        "    .muted {{ color: #6e6e73; }}",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>beaver-gateway</h1>",
        '  <p class="muted">Signed in as <strong>{actor}</strong>.'
        " Base URL: <code>{base_url}</code></p>",
        "  <h2>Namespaces</h2>",
        "  <table>",
        "    <thead><tr><th>Name</th><th>URL</th></tr></thead>",
        "    <tbody>",
        "{rows}",
        "    </tbody>",
        "  </table>",
        '  <p class="muted">Bundle endpoint (flat namespace,'
        " escape-hatch): <code>{all_url}</code></p>",
        "  <h2>Cursor / Cline</h2>",
        "  <p>Add to your MCP config:</p>",
        "  <pre>{cursor_snippet}</pre>",
        "  <h2>Claude Desktop</h2>",
        "  <p>Add to <code>"
        "~/Library/Application Support/Claude/claude_desktop_config.json"
        "</code>:</p>",
        "  <pre>{claude_desktop_snippet}</pre>",
        "</body>",
        "</html>",
        "",
    ]
)


_CURSOR_SNIPPET = """{{
  "mcpServers": {{
    "beaver-time": {{
      "url": "{base_url}/time/",
      "headers": {{ "Authorization": "Bearer <YOUR_TOKEN>" }}
    }}
  }}
}}"""


_CLAUDE_DESKTOP_SNIPPET = """{{
  "mcpServers": {{
    "beaver-time": {{
      "type": "http",
      "url": "{base_url}/time/",
      "headers": {{ "Authorization": "Bearer <YOUR_TOKEN>" }}
    }}
  }}
}}"""