from typing import Any

import httpx
import jwt

from app.config import settings

_jwks_cache: dict[str, Any] | None = None


async def decode_token(token: str) -> dict[str, Any]:
    if settings.authentik_issuer_url:
        issuer = settings.authentik_issuer_url.rstrip("/")
        jwks = await _get_jwks(issuer)

        signing_key = jwt.algorithms.RSAAlgorithm.from_jwk(
            _find_matching_key(jwks, token)
        )

        return jwt.decode(
            token,
            signing_key,  # type: ignore[arg-type]
            algorithms=["RS256"],
            audience=settings.authentik_client_id,
            issuer=settings.authentik_issuer_url,
        )

    return jwt.decode(token, options={"verify_signature": False})


async def _get_jwks(issuer: str) -> dict[str, Any]:
    global _jwks_cache

    if _jwks_cache is not None:
        return _jwks_cache

    discovery_url = f"{issuer}/.well-known/openid-configuration"
    async with httpx.AsyncClient() as client:
        resp = await client.get(discovery_url)
        resp.raise_for_status()
        discovery = resp.json()
        jwks_uri = discovery["jwks_uri"]

        jwks_resp = await client.get(jwks_uri)
        jwks_resp.raise_for_status()
        _jwks_cache = jwks_resp.json()

    return _jwks_cache


def _find_matching_key(jwks: dict[str, Any], token: str) -> dict[str, Any]:
    unverified_header = jwt.get_unverified_header(token)
    kid = unverified_header.get("kid")
    for key in jwks.get("keys", []):
        key_dict: dict[str, Any] = key
        if key_dict.get("kid") == kid:
            return key_dict
    raise RuntimeError(f"No matching JWKS key found for kid={kid}")