# Copyright (C) 2026 Garudex Labs.  All Rights Reserved.
# Caracal, a product of Garudex Labs
#
# JWKS cache with 5-min TTL.

import asyncio
import time
import httpx
from caracalai_core import JsonValue

_TTL = 300.0


class JwksCache:
    def __init__(self) -> None:
        self._cache: dict[str, tuple[list[dict[str, JsonValue]], float]] = {}
        self._locks: dict[str, asyncio.Lock] = {}
        self._locks_guard = asyncio.Lock()

    async def _lock_for(self, issuer: str) -> asyncio.Lock:
        async with self._locks_guard:
            lock = self._locks.get(issuer)
            if lock is None:
                lock = asyncio.Lock()
                self._locks[issuer] = lock
            return lock

    async def get_keys(self, issuer: str) -> list[dict[str, JsonValue]]:
        url = issuer.rstrip("/") + "/.well-known/jwks.json"
        entry = self._cache.get(issuer)
        if entry and time.monotonic() - entry[1] < _TTL:
            return entry[0]

        # Per-issuer lock coalesces concurrent fetches: the second caller
        # waits, then reads the freshly-cached entry instead of re-fetching.
        lock = await self._lock_for(issuer)
        async with lock:
            entry = self._cache.get(issuer)
            if entry and time.monotonic() - entry[1] < _TTL:
                return entry[0]

            async with httpx.AsyncClient() as client:
                resp = await client.get(url)
                resp.raise_for_status()
                body = resp.json()

            keys = body.get("keys", []) if isinstance(body, dict) else []
            keys = [key for key in keys if isinstance(key, dict)]
            self._cache[issuer] = (keys, time.monotonic())
            return keys
