"""JWKS client with TTL cache.

Stateless: downloads the realm public keys and caches them. Token validation does
NOT call Keycloak on every request — it only refreshes the JWKS when the TTL expires
or an unknown `kid` shows up (key rotation).
"""

import time

import httpx
import jwt
from jwt import PyJWK


class JWKSCache:
    def __init__(self, jwks_uri: str, ttl: int) -> None:
        self._uri = jwks_uri
        self._ttl = ttl
        self._jwks: dict | None = None
        self._fetched_at: float = 0.0

    async def _refresh(self) -> None:
        async with httpx.AsyncClient(timeout=5.0) as client:
            resp = await client.get(self._uri)
            resp.raise_for_status()
            self._jwks = resp.json()
            self._fetched_at = time.monotonic()

    async def get_key(self, kid: str) -> PyJWK:
        # refresh if the cache expired or the kid is missing (possible rotation)
        expired = (time.monotonic() - self._fetched_at) > self._ttl
        if self._jwks is None or expired or not self._has_kid(kid):
            await self._refresh()
        assert self._jwks is not None
        for key in self._jwks["keys"]:
            if key["kid"] == kid:
                return PyJWK.from_dict(key)
        raise jwt.PyJWKError(f"kid {kid!r} not found in JWKS")

    def _has_kid(self, kid: str) -> bool:
        if self._jwks is None:
            return False
        return any(k["kid"] == kid for k in self._jwks["keys"])
