"""Workaround for FastMCP issue #3509: refresh_expires_in=0 from Keycloak offline tokens.

Keycloak returns ``refresh_expires_in=0`` for offline_access tokens to indicate
"no fixed time-based expiry — validity is governed by realm offline session settings".
FastMCP treats 0 as falsy and skips issuing PROXY_RT entirely, so MCP clients must
re-authenticate in the browser after every access token expiry (typically 1 hour).

The mixin patches the stored IdP token response before FastMCP's
``exchange_authorization_code`` processes it, replacing 0 with
``PROXY_SESSION_TTL_SECONDS`` (30 days) so FastMCP correctly issues and stores
the refresh token.

Remove once https://github.com/PrefectHQ/fastmcp/issues/3509 is resolved and
we upgrade to a fixed FastMCP version.
"""

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)

# Short TTL for the re-stored code entry: super() deletes it immediately after
# fetching, so anything > a few seconds is sufficient.
_REPUT_TTL_SECONDS = 30


def with_offline_refresh_fix(base_cls: type) -> type:
    """Return a subclass of *base_cls* with the offline refresh fix applied.

    Usage::

        OIDCProxy = with_offline_refresh_fix(OIDCProxy)
        OAuthProxy = with_offline_refresh_fix(OAuthProxy)
    """
    return type(base_cls.__name__, (OfflineRefreshFixMixin, base_cls), {})


class OfflineRefreshFixMixin:
    """Mixin for OAuthProxy / OIDCProxy that fixes refresh_expires_in=0.

    Apply before the FastMCP base class in the MRO:

        class _PatchedOIDCProxy(OfflineRefreshFixMixin, OIDCProxy): ...
        class _PatchedOAuthProxy(OfflineRefreshFixMixin, OAuthProxy): ...
    """

    async def exchange_authorization_code(
        self,
        client: Any,
        authorization_code: Any,
    ) -> Any:
        from ims_mcp.constants import PROXY_SESSION_TTL_SECONDS

        code_model = await self._code_store.get(key=authorization_code.code)  # type: ignore[attr-defined]
        if (
            code_model is not None
            and code_model.idp_tokens.get("refresh_token")
            and code_model.idp_tokens.get("refresh_expires_in") == 0
        ):
            # 0 = Keycloak offline token convention: "no fixed expiry".
            # FastMCP treats 0 as falsy and never issues PROXY_RT.
            # Patch to 30 days so FastMCP issues and stores the refresh token.
            logger.debug(
                "offline_access fix: patching refresh_expires_in 0 → %d s "
                "(upstream is a Keycloak offline token; see fastmcp#3509)",
                PROXY_SESSION_TTL_SECONDS,
            )
            code_model.idp_tokens["refresh_expires_in"] = PROXY_SESSION_TTL_SECONDS
            await self._code_store.put(  # type: ignore[attr-defined]
                key=authorization_code.code,
                value=code_model,
                ttl=_REPUT_TTL_SECONDS,
            )

        return await super().exchange_authorization_code(client, authorization_code)
