"""SSE transport + bearer auth for the MCP server.

Mounts two routes on the live FastAPI server:

- ``GET  /mcp/sse``      — opens the SSE stream. Client holds this open.
- ``POST /mcp/messages/`` — receives the client's outgoing messages, correlates
                           them to a session established on the SSE stream.

Both routes require ``Authorization: Bearer <raw_token>`` (or
``?token=<raw_token>`` for ``EventSource`` clients that can't set headers —
browsers).

**Auth is fail-closed.** If the token store is empty (no tokens minted yet),
every request returns 401 — we never expose tool execution on an open port.
Cross-device MCP is opt-in, not default-on.

The SSE server reuses ``app.state.tool_executor`` — the SAME executor the
chat UI uses. That means every tool call inherits contract enforcement,
activity logging, decision records, trust gating, and the correlation
scheme. No duplicate state, no second copy of the graph.
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING

from fastapi import APIRouter, HTTPException, Request
from mcp.server.sse import SseServerTransport
from starlette.responses import Response

if TYPE_CHECKING:
    from agntspace.mcp.sessions import SseSessionTracker
    from agntspace.mcp.tokens import TokenStore

log = logging.getLogger(__name__)


_MESSAGES_PATH = "/mcp/messages/"


def _extract_token(request: Request) -> str | None:
    """Pull the bearer token from ``Authorization`` header or ``?token=`` query.

    Header takes precedence. The query fallback exists because browser
    ``EventSource`` can't set custom headers — some MCP-over-browser
    clients route through it.
    """
    auth = request.headers.get("authorization") or request.headers.get("Authorization")
    if auth:
        parts = auth.split(" ", 1)
        if len(parts) == 2 and parts[0].lower() == "bearer":
            return parts[1].strip() or None
    return request.query_params.get("token") or None


async def _authenticate(request: Request, tokens: TokenStore):
    """Return the verified token (id + scope) on success; raise 401 on failure.

    Never leak WHY auth failed — "invalid token" covers missing, malformed,
    revoked, and wrong-secret identically. Defense against account enumeration.

    The returned ``VerifiedToken`` carries the scope the bearer was minted
    under; callers use that to shape the ``Server`` instance so scope
    enforcement happens per-session, not globally.
    """
    if not await tokens.any_live():
        # No tokens minted → SSE disabled. Clients get a clear hint via
        # the response body, not a 200 with empty content.
        raise HTTPException(
            status_code=401,
            detail=(
                "MCP SSE is disabled — no tokens configured. "
                "Create one: POST /api/mcp/tokens"
            ),
        )
    raw = _extract_token(request)
    if not raw:
        raise HTTPException(status_code=401, detail="Missing bearer token")
    verified = await tokens.verify(raw)
    if not verified:
        raise HTTPException(status_code=401, detail="Invalid token")
    return verified


async def _enforce_rate(scope: str, token_id: str) -> None:
    """Per-token, per-scope rate gate. Raises 429 when the bucket is empty.

    The limiter is a process-wide singleton; ``caller=token_id`` keeps
    each token in its own bucket so a single noisy client doesn't starve
    the rest. Scopes + capacities live in ``infra.rate_limit.get_limiter``.

    Failures here are NEVER swallowed — exceeding the limit is a user
    signal, not an error to log and continue past.
    """
    from agntspace.infra.rate_limit import get_limiter
    limiter = get_limiter()
    if not await limiter.acquire(scope, caller=token_id):
        raise HTTPException(
            status_code=429,
            detail=f"Rate limit exceeded for {scope}. Slow down or revoke + re-mint a fresh token.",
        )


def build_sse_router(transport: SseServerTransport) -> APIRouter:
    """Build the FastAPI router exposing the two SSE endpoints.

    Both routes read ``request.app.state`` for ``tool_executor``, ``vault``,
    and ``mcp_tokens`` — wiring done in ``main.py``. That indirection keeps
    this module import-light and testable.
    """
    router = APIRouter()

    @router.get("/mcp/sse")
    async def handle_sse(request: Request) -> Response:
        from agntspace.mcp.server import create_mcp_server

        tokens: TokenStore = request.app.state.mcp_tokens
        verified = await _authenticate(request, tokens)
        # Per-token rate gate on session opens. Capped at 10/min so a
        # broken client can't reconnect-storm the server.
        await _enforce_rate("mcp_session_open", verified.token_id)

        tool_executor = getattr(request.app.state, "tool_executor", None)
        if tool_executor is None:
            # Server booted without the full stack (unusual — setup mode?).
            raise HTTPException(
                status_code=503,
                detail="Tool executor not available on this server instance",
            )
        vault = getattr(request.app.state, "vault", None)
        vault_path = vault.vault_dir if vault else None

        # Per-session server. Scope travels with the token — a read_only
        # token gets a Server whose list_tools and call_tool both enforce
        # the narrowing. Full-scope tokens get the full surface.
        server = create_mcp_server(
            tool_executor,
            vault_path=vault_path,
            vault_reader=vault,
            scope=verified.scope,
        )

        # Register the session so the revoke API can force-close us.
        # ``tracker`` is set up in main.py; if it's missing (e.g., very
        # early boot or a setup-mode app) we degrade gracefully — the
        # session still runs, it just can't be force-closed mid-stream.
        tracker: SseSessionTracker | None = getattr(request.app.state, "mcp_sse_tracker", None)
        handle = await tracker.register(verified.token_id) if tracker else None

        log.info(
            "MCP SSE session opened (token=%s, scope=%s)",
            verified.token_id, verified.scope,
        )
        try:
            async with transport.connect_sse(
                request.scope, request.receive, request._send,  # noqa: SLF001 — starlette pattern
            ) as (read_stream, write_stream):
                # Race the MCP server loop against the revoke signal.
                # Whichever finishes first wins; the loser is cancelled
                # and awaited so it can clean up its streams. Without
                # this, a revoked-but-still-streaming token would keep
                # serving until the client disconnects.
                init_opts = server.create_initialization_options()
                server_task = asyncio.create_task(
                    server.run(read_stream, write_stream, init_opts),
                    name=f"mcp-sse-{verified.token_id}",
                )
                if handle is None:
                    # No tracker available — just await the server normally.
                    waiters = {server_task}
                else:
                    revoke_task = asyncio.create_task(
                        handle.revoke_event.wait(),
                        name=f"mcp-sse-revoke-{verified.token_id}",
                    )
                    waiters = {server_task, revoke_task}

                done, pending = await asyncio.wait(
                    waiters, return_when=asyncio.FIRST_COMPLETED,
                )

                # Cancel whatever didn't win the race and drain it
                # so we don't leak coroutines or transport resources.
                for t in pending:
                    t.cancel()
                    try:
                        await t
                    except (asyncio.CancelledError, Exception):
                        pass

                if handle is not None and handle.revoke_event.is_set():
                    log.info(
                        "MCP SSE session force-closed on token revoke (token=%s)",
                        verified.token_id,
                    )
                # Surface server-task exceptions that finished naturally
                # (i.e., not cancelled by revoke) so the outer except logs them.
                if server_task in done:
                    exc = server_task.exception()
                    if exc is not None:
                        raise exc
        except Exception:
            log.exception("MCP SSE session errored (token=%s)", verified.token_id)
            raise
        finally:
            if tracker is not None and handle is not None:
                await tracker.unregister(handle)
            log.info("MCP SSE session closed (token=%s)", verified.token_id)

        # Must return a Response or Starlette raises on client-disconnect.
        return Response()

    @router.post(_MESSAGES_PATH)
    async def handle_messages(request: Request) -> Response:
        tokens: TokenStore = request.app.state.mcp_tokens
        verified = await _authenticate(request, tokens)
        # Per-token, per-scope message rate gate. Read-only tokens get
        # a more generous bucket since the blast radius is just CPU +
        # token cost, not data mutation.
        rate_scope = (
            "mcp_message_read_only"
            if verified.scope == "read_only"
            else "mcp_message_full"
        )
        await _enforce_rate(rate_scope, verified.token_id)
        # ``handle_post_message`` is itself an ASGI application — we hand
        # it the same scope/receive/send triplet FastAPI gave us and let
        # it write the response. No body wrangling on our side.
        await transport.handle_post_message(request.scope, request.receive, request._send)  # noqa: SLF001
        # The ASGI app already sent the response; return an empty one so
        # Starlette doesn't complain about the missing return.
        return Response()

    return router


def create_transport() -> SseServerTransport:
    """Factory — keeps the path constant in one place."""
    return SseServerTransport(_MESSAGES_PATH)
