"""JWT token validation against Microsoft Entra ID JWKS keys."""

import re
import time
import threading

import httpx
import jwt
from jwt.algorithms import RSAAlgorithm

from .config import AuthSettings
from .exceptions import AuthenticationError, TokenExpiredError
from .models import AuthenticatedUser

# Multi-tenant issuer patterns (v2.0 and v1.0 endpoints)
_ISSUER_V2 = re.compile(
    r"https://login\.microsoftonline\.com/"
    r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/v2\.0"
)
_ISSUER_V1 = re.compile(
    r"https://sts\.windows\.net/"
    r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/"
)


class JWKSKeyCache:
    """Thread-safe cache for Microsoft's JWKS signing keys."""

    def __init__(self) -> None:
        self._keys: list[dict] = []
        self._fetched_at: float = 0
        self._lock = threading.Lock()

    def get_keys(self, settings: AuthSettings) -> list[dict]:
        """Get JWKS keys, fetching from Microsoft if cache has expired."""
        now = time.time()
        if now - self._fetched_at < settings.JWKS_CACHE_TTL and self._keys:
            return self._keys

        with self._lock:
            # Double-check after acquiring lock
            if now - self._fetched_at < settings.JWKS_CACHE_TTL and self._keys:
                return self._keys
            return self._fetch_sync(settings)

    async def get_keys_async(self, settings: AuthSettings) -> list[dict]:
        """Get JWKS keys asynchronously."""
        now = time.time()
        if now - self._fetched_at < settings.JWKS_CACHE_TTL and self._keys:
            return self._keys

        return await self._fetch_async(settings)

    def _fetch_sync(self, settings: AuthSettings) -> list[dict]:
        resp = httpx.get(settings.JWKS_URI, timeout=10)
        resp.raise_for_status()
        self._keys = resp.json()["keys"]
        self._fetched_at = time.time()
        return self._keys

    async def _fetch_async(self, settings: AuthSettings) -> list[dict]:
        async with httpx.AsyncClient() as client:
            resp = await client.get(settings.JWKS_URI, timeout=10)
            resp.raise_for_status()
            self._keys = resp.json()["keys"]
            self._fetched_at = time.time()
            return self._keys

    def clear(self) -> None:
        """Clear the cache (useful for testing)."""
        self._keys = []
        self._fetched_at = 0


# Module-level cache instance
_key_cache = JWKSKeyCache()


def get_key_cache() -> JWKSKeyCache:
    """Get the module-level JWKS key cache."""
    return _key_cache


def validate_token(
    token: str,
    jwks_keys: list[dict],
    settings: AuthSettings,
) -> AuthenticatedUser:
    """Validate a JWT and return an AuthenticatedUser.

    Args:
        token: The raw JWT string (without "Bearer " prefix).
        jwks_keys: List of JWKS key dicts from Microsoft's endpoint.
        settings: Auth configuration.

    Returns:
        AuthenticatedUser with claims extracted from the token.

    Raises:
        AuthenticationError: If the token is invalid.
        TokenExpiredError: If the token has expired.
    """
    audience = settings.ENTRA_AUDIENCE or settings.ENTRA_CLIENT_ID

    # Decode header to find the signing key
    try:
        unverified_header = jwt.get_unverified_header(token)
    except jwt.exceptions.DecodeError as e:
        raise AuthenticationError(f"Malformed token header: {e}")

    kid = unverified_header.get("kid")
    if not kid:
        raise AuthenticationError("Token header missing 'kid' field")

    matching_key = next((k for k in jwks_keys if k["kid"] == kid), None)
    if not matching_key:
        raise AuthenticationError("Token signing key not found in JWKS")

    # Build the public key from the JWK
    try:
        public_key = RSAAlgorithm.from_jwk(matching_key)
    except Exception as e:
        raise AuthenticationError(f"Failed to construct public key: {e}")

    # Decode and validate the token
    try:
        claims = jwt.decode(
            token,
            public_key,
            algorithms=["RS256"],
            audience=audience,
            options={"verify_iss": False},  # We validate issuer manually below
        )
    except jwt.ExpiredSignatureError:
        raise TokenExpiredError("Token has expired")
    except jwt.InvalidAudienceError:
        raise AuthenticationError(
            f"Token audience mismatch: expected '{audience}'"
        )
    except jwt.InvalidTokenError as e:
        raise AuthenticationError(f"Invalid token: {e}")

    # Multi-tenant issuer validation: accept any valid Azure AD tenant
    issuer = claims.get("iss", "")
    if not (_ISSUER_V2.match(issuer) or _ISSUER_V1.match(issuer)):
        raise AuthenticationError(f"Unexpected token issuer: {issuer}")

    return AuthenticatedUser(
        tenant_id=claims.get("tid", ""),
        user_id=claims.get("oid", ""),
        email=claims.get("preferred_username", ""),
        display_name=claims.get("name", ""),
        roles=claims.get("roles", []),
        raw_claims=claims,
    )
