app/utils/auth.py

import jwt
import httpx
import time
from jwt.algorithms import ECAlgorithm
from app.config import settings

# ---------------------------------------------------------------------------
# JWKS cache — fetched once and refreshed every 12 hours
# ---------------------------------------------------------------------------
_jwks_cache: dict = {}
_jwks_cache_ts: float = 0.0
_JWKS_TTL = 12 * 60 * 60  # 12 hours in seconds

SUPABASE_JWKS_URL = f"{settings.SUPABASE_URL}/auth/v1/.well-known/jwks.json"


def _fetch_jwks_sync() -> dict:
    """Synchronous JWKS fetch — called via run_in_executor from async context."""
    response = httpx.get(SUPABASE_JWKS_URL, timeout=5.0)
    response.raise_for_status()
    return response.json()


async def _get_jwks_async() -> dict:
    """Fetch JWKS from Supabase with in-process cache. Non-blocking for async event loop."""
    global _jwks_cache, _jwks_cache_ts
    now = time.time()
    if _jwks_cache and (now - _jwks_cache_ts) < _JWKS_TTL:
        return _jwks_cache

    import asyncio
    try:
        loop = asyncio.get_event_loop()
        data = await loop.run_in_executor(None, _fetch_jwks_sync)
        _jwks_cache = {key["kid"]: key for key in data.get("keys", [])}
        _jwks_cache_ts = now
        return _jwks_cache
    except Exception as e:
        if _jwks_cache:  # Use stale cache on failure rather than crash
            return _jwks_cache
        raise RuntimeError(f"Failed to fetch Supabase JWKS: {e}") from e


def _get_jwks() -> dict:
    """Synchronous fallback — used only by non-async callers."""
    global _jwks_cache, _jwks_cache_ts
    now = time.time()
    if _jwks_cache and (now - _jwks_cache_ts) < _JWKS_TTL:
        return _jwks_cache
    try:
        response = httpx.get(SUPABASE_JWKS_URL, timeout=5.0)
        response.raise_for_status()
        data = response.json()
        _jwks_cache = {key["kid"]: key for key in data.get("keys", [])}
        _jwks_cache_ts = now
        return _jwks_cache
    except Exception as e:
        if _jwks_cache:
            return _jwks_cache
        raise RuntimeError(f"Failed to fetch Supabase JWKS: {e}") from e


def _get_public_key(kid: str):
    """Return the PyJWT-compatible public key object for the given key ID."""
    jwks = _get_jwks()
    jwk = jwks.get(kid)
    if not jwk:
        # Refresh cache once and retry (key rotation scenario)
        global _jwks_cache_ts
        _jwks_cache_ts = 0.0
        jwks = _get_jwks()
        jwk = jwks.get(kid)
    if not jwk:
        raise ValueError(f"Unknown JWKS key ID: {kid}")
    return ECAlgorithm.from_jwk(jwk)


def verify_supabase_jwt(token: str) -> dict:
    """
    Decodes and verifies a Supabase-issued JWT.

    Supports both:
    - ES256 (ECC P-256)  — current Supabase signing algorithm via JWKS
    - HS256              — legacy shared-secret tokens (fallback)

    Returns a normalised claims dict compatible with the rest of the backend.
    """
    # Decode header without verification to determine the algorithm and kid
    try:
        unverified_header = jwt.get_unverified_header(token)
    except jwt.exceptions.DecodeError as e:
        raise jwt.DecodeError(f"Malformed JWT header: {e}") from e

    alg = unverified_header.get("alg", "")
    kid = unverified_header.get("kid")

    if alg == "ES256" and kid:
        # ----------------------------------------------------------------
        # Verify using JWKS (ECC P-256 — current Supabase default)
        # ----------------------------------------------------------------
        public_key = _get_public_key(kid)
        try:
            claims = jwt.decode(
                token,
                public_key,
                algorithms=["ES256"],
                options={"verify_aud": True},
                audience="authenticated",
            )
        except jwt.InvalidAudienceError:
            claims = jwt.decode(
                token,
                public_key,
                algorithms=["ES256"],
                options={"verify_aud": False},
            )
    else:
        # ----------------------------------------------------------------
        # Fallback: HS256 verification with shared JWT secret
        # ----------------------------------------------------------------
        try:
            claims = jwt.decode(
                token,
                settings.SUPABASE_JWT_SECRET,
                algorithms=["HS256"],
                options={"verify_aud": True},
                audience="authenticated",
            )
        except jwt.InvalidAudienceError:
            claims = jwt.decode(
                token,
                settings.SUPABASE_JWT_SECRET,
                algorithms=["HS256"],
                options={"verify_aud": False},
            )

    # ------------------------------------------------------------------
    # Normalise claims to a consistent dict expected by the rest of the app
    # ------------------------------------------------------------------
    user_metadata = claims.get("user_metadata", {}) or {}
    uid = claims.get("sub")
    email = claims.get("email")
    name = (
        user_metadata.get("full_name")
        or claims.get("name")
        or user_metadata.get("name")
    )
    picture = (
        user_metadata.get("avatar_url")
        or claims.get("picture")
        or user_metadata.get("picture")
    )

    if not uid:
        raise ValueError("Missing sub (subject) claim in JWT")
    if not email:
        raise ValueError("Missing email claim in JWT")

    return {
        "uid": uid,
        "sub": uid,
        "email": email,
        "user_metadata": user_metadata,
        "full_name": name,
        "name": name,
        "avatar_url": picture,
        "picture": picture,
        "exp": claims.get("exp"),
    }