"""Token exchange: the end user's token in, one narrowly-scoped token per MCP
server out.

Two measured facts about Keycloak 26.6 shape this, both of which cost time to
discover and neither of which is obvious from the RFC:

**`actor_token` is silently ignored.** Standard Token Exchange accepts the
parameter, returns a valid token, and emits no `act` claim — no error, no
warning. Delegation therefore cannot be built on RFC 8693 `act`. The actor is
identified by `azp`, which carries this agent's client id.

**The audience is requested as a scope.** Passing `audience=mcp-catalog` returns
"Requested audience not available". Passing `scope=openid mcp-catalog-audience`
works, and the resulting token carries `aud: [mcp-catalog]`.
"""

from __future__ import annotations

import base64
import json
import time
from dataclasses import dataclass

import httpx2

GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"


class TokenExchangeError(RuntimeError):
    """Keycloak refused the exchange. Carries its response, which names the cause."""


@dataclass(frozen=True)
class ExchangedToken:
    access_token: str
    expires_at: float

    @property
    def expired(self) -> bool:
        # A conversation turn lasts seconds and a token five minutes, so this
        # margin only matters for a turn that stalls on a slow model.
        return time.monotonic() >= self.expires_at - 30


class TokenExchanger:
    """Exchanges a user token for audience-scoped tokens.

    Caches per (subject token, scope) for the lifetime of one turn. Deliberately
    not a long-lived cache keyed by user: a token that outlives the turn it was
    minted for is a token that can be replayed after the user's session ended.
    """

    def __init__(self, client: httpx2.AsyncClient, *, token_endpoint: str, client_id: str, client_secret: str) -> None:
        self._client = client
        self._endpoint = token_endpoint
        self._client_id = client_id
        self._client_secret = client_secret
        self._cache: dict[tuple[str, str], ExchangedToken] = {}

    async def exchange(self, subject_token: str, audience_scope: str) -> str:
        key = (subject_token, audience_scope)
        cached = self._cache.get(key)
        if cached and not cached.expired:
            return cached.access_token

        response = await self._client.post(
            self._endpoint,
            data={
                "grant_type": GRANT_TYPE,
                "client_id": self._client_id,
                "client_secret": self._client_secret,
                "subject_token": subject_token,
                "subject_token_type": SUBJECT_TOKEN_TYPE,
                # Not `audience=...`. See the module docstring.
                "scope": f"openid {audience_scope}",
            },
        )
        if response.status_code != 200:
            raise TokenExchangeError(
                f"exchange for {audience_scope!r} failed: HTTP {response.status_code} {response.text[:200]}"
            )

        payload = response.json()
        token = ExchangedToken(
            access_token=payload["access_token"],
            expires_at=time.monotonic() + float(payload.get("expires_in", 60)),
        )
        self._cache[key] = token
        return token.access_token


def caller_of(access_token: str) -> str:
    """The username inside a token, for the prompt and nothing else.

    Read without verifying the signature on purpose: this is a label shown to
    the model, never a decision. Every decision that depends on this identity is
    taken by Keycloak or by the gateway, both of which do verify.
    """
    try:
        payload = access_token.split(".")[1]
        payload += "=" * (-len(payload) % 4)
        claims = json.loads(base64.urlsafe_b64decode(payload))
    except Exception:  # noqa: BLE001 — a malformed token is the gateway's problem
        return "an unidentified caller"
    return claims.get("preferred_username") or claims.get("sub") or "an unidentified caller"