"""OAuthProxy subclass that transparently refreshes upstream tokens.

MCP clients (Claude Code, Cursor) may not implement the OAuth refresh
token flow.  This proxy refreshes the upstream Keycloak access token
server-side when introspection returns ``active=false``, so the client
never sees ``invalid_token`` as long as the upstream refresh token is
still alive.

Two adjustments over vanilla ``OAuthProxy``:

1. **Lenient JWT verification** — ``load_access_token`` checks the
   FastMCP JWT signature, issuer, and audience but *skips* the ``exp``
   claim.  The JWT is a reference token; real authorization comes from
   upstream introspection.

2. **Transparent upstream refresh** — when introspection fails but a
   refresh token is stored, the proxy calls Keycloak's token endpoint to
   obtain a fresh access token, updates the store, and re-introspects.

3. **Extended JTI mapping TTL** — after ``exchange_authorization_code``
   and ``exchange_refresh_token``, the JTI mapping TTL is extended to
   match the refresh token lifetime so it outlives the short-lived
   access token.
"""

from __future__ import annotations

import logging
import time
from typing import Any

import httpx
from mcp.server.auth.provider import AccessToken

from ims_mcp.constants import PROXY_SESSION_TTL_SECONDS

logger = logging.getLogger(__name__)


class RefreshingOAuthProxyMixin:
    """Mixin that adds transparent upstream token refresh to OAuthProxy.

    Applied via multiple inheritance so the concrete class is:
        ``class RefreshingOAuthProxy(RefreshingOAuthProxyMixin, OAuthProxy)``
    MRO ensures ``super()`` calls hit ``OAuthProxy`` methods.
    """

    # ------------------------------------------------------------------
    # Lenient JWT decode (skip ``exp``)
    # ------------------------------------------------------------------

    def _decode_jwt_lenient(self, token: str) -> dict[str, Any] | None:
        """Decode a FastMCP JWT checking signature, issuer, audience only.

        Skips the ``exp`` claim so that expired-but-signed tokens are
        still accepted as valid reference tokens.
        """
        try:
            payload = self.jwt_issuer._jwt.decode(  # type: ignore[attr-defined]
                token,
                self.jwt_issuer._signing_key,  # type: ignore[attr-defined]
            )
            if payload.get("iss") != self.jwt_issuer.issuer:  # type: ignore[attr-defined]
                return None
            if payload.get("aud") != self.jwt_issuer.audience:  # type: ignore[attr-defined]
                return None
            return dict(payload)
        except Exception:
            return None

    # ------------------------------------------------------------------
    # Upstream refresh helper
    # ------------------------------------------------------------------

    async def _refresh_upstream_directly(
        self, refresh_token: str
    ) -> dict[str, Any] | None:
        """Call Keycloak's token endpoint with the stored refresh token."""
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                response = await client.post(
                    self._upstream_token_endpoint,  # type: ignore[attr-defined]
                    data={
                        "grant_type": "refresh_token",
                        "refresh_token": refresh_token,
                        "client_id": self._upstream_client_id,  # type: ignore[attr-defined]
                        "client_secret": self._upstream_client_secret.get_secret_value(),  # type: ignore[attr-defined]
                    },
                )
                if response.status_code != 200:
                    logger.debug(
                        "Upstream refresh failed: HTTP %d", response.status_code
                    )
                    return None
                return response.json()  # type: ignore[no-any-return]
        except Exception as exc:
            logger.debug("Upstream refresh error: %s", exc)
            return None

    # ------------------------------------------------------------------
    # Overridden load_access_token with transparent refresh
    # ------------------------------------------------------------------

    async def load_access_token(self, token: str) -> AccessToken | None:
        """Validate FastMCP JWT, introspect upstream token, refresh if needed."""
        # 1. Decode JWT leniently (skip exp)
        payload = self._decode_jwt_lenient(token)
        if payload is None:
            return None
        jti: str | None = payload.get("jti")
        if not jti:
            return None

        # 2. Look up JTI → upstream token mapping
        jti_mapping = await self._jti_mapping_store.get(key=jti)  # type: ignore[attr-defined]
        if not jti_mapping:
            logger.debug("JTI mapping expired or not found: %s", jti[:16])
            return None

        upstream = await self._upstream_token_store.get(  # type: ignore[attr-defined]
            key=jti_mapping.upstream_token_id
        )
        if not upstream:
            return None

        # 3. Introspect the current upstream access token
        verification_token = self._get_verification_token(upstream)  # type: ignore[attr-defined]
        validated: AccessToken | None = None
        if verification_token:
            validated = await self._token_validator.verify_token(verification_token)  # type: ignore[attr-defined]

        # 4. If introspection failed → try transparent refresh
        if validated is None and upstream.refresh_token:
            logger.debug("Upstream token expired, attempting transparent refresh")
            new_tokens = await self._refresh_upstream_directly(upstream.refresh_token)
            if new_tokens and "access_token" in new_tokens:
                # Update upstream token set
                upstream.access_token = new_tokens["access_token"]
                new_expires_in = int(new_tokens.get("expires_in", 300))
                upstream.expires_at = time.time() + new_expires_in

                if "refresh_token" in new_tokens:
                    upstream.refresh_token = new_tokens["refresh_token"]
                if "refresh_expires_in" in new_tokens:
                    upstream.refresh_token_expires_at = time.time() + int(
                        new_tokens["refresh_expires_in"]
                    )

                # Persist updated upstream tokens
                refresh_ttl = (
                    int(upstream.refresh_token_expires_at - time.time())
                    if upstream.refresh_token_expires_at
                    else PROXY_SESSION_TTL_SECONDS
                )
                await self._upstream_token_store.put(  # type: ignore[attr-defined]
                    key=jti_mapping.upstream_token_id,
                    value=upstream,
                    ttl=max(refresh_ttl, new_expires_in, 1),
                )

                # Re-introspect with new access token
                verification_token = self._get_verification_token(upstream)  # type: ignore[attr-defined]
                if verification_token:
                    validated = await self._token_validator.verify_token(  # type: ignore[attr-defined]
                        verification_token
                    )
                    if validated:
                        logger.debug("Transparent upstream refresh succeeded")

        if validated is None:
            return None

        # 5. Fix up token fields (same logic as parent)
        if verification_token and verification_token != upstream.access_token:
            validated = validated.model_copy(
                update={
                    "token": upstream.access_token,
                    "scopes": upstream.scope.split() if upstream.scope else validated.scopes,
                    "expires_at": int(upstream.expires_at),
                }
            )

        return validated

    # ------------------------------------------------------------------
    # Extend JTI mapping TTL after token issuance
    # ------------------------------------------------------------------

    async def _extend_jti_ttl(self, access_token_jwt: str) -> None:
        """Re-store the JTI mapping with a longer TTL."""
        payload = self._decode_jwt_lenient(access_token_jwt)
        if not payload:
            return
        jti = payload.get("jti")
        if not jti:
            return
        mapping = await self._jti_mapping_store.get(key=jti)  # type: ignore[attr-defined]
        if mapping:
            await self._jti_mapping_store.put(  # type: ignore[attr-defined]
                key=jti,
                value=mapping,
                ttl=PROXY_SESSION_TTL_SECONDS,
            )

    async def exchange_authorization_code(self, client: Any, authorization_code: Any) -> Any:
        """Delegate to parent, then extend JTI mapping TTL."""
        result = await super().exchange_authorization_code(client, authorization_code)
        await self._extend_jti_ttl(result.access_token)
        return result

    async def exchange_refresh_token(self, client: Any, refresh_token: Any, scopes: Any) -> Any:
        """Delegate to parent, then extend JTI mapping TTL."""
        result = await super().exchange_refresh_token(client, refresh_token, scopes)
        await self._extend_jti_ttl(result.access_token)
        return result
