from typing import Any

import jwt

from app.config import settings


def decode_token(token: str) -> dict[str, Any]:
    """Decode a JWT token.

    When authentik_issuer_url is configured, validates the token
    against the OIDC discovery document JWKS.
    Otherwise, decodes without verification (local development only).
    """
    if settings.authentik_issuer_url:
        import httpx

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

            jwks_resp = client.get(jwks_uri)
            jwks_resp.raise_for_status()
            jwks = jwks_resp.json()

        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})


def _find_matching_key(jwks: dict[str, Any], token: str) -> dict[str, Any]:
    """Find the key in JWKS that matches the token's kid header."""
    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}")