"""JWKS-backed bearer-token verification.

The Node SDK ships an Express middleware; frameworks differ too much in Python
for a single drop-in middleware, so the core is exposed as a framework-agnostic
:class:`TokenVerifier` plus a :func:`extract_bearer_token` helper. Wire these
into your framework's middleware/dependency of choice (Flask, FastAPI, ...).
"""

from __future__ import annotations

from typing import Any, Mapping, Optional, Sequence, Union

import jwt
from jwt import PyJWKClient

from .config import ResolvedConfig

Audience = Union[str, Sequence[str]]


class TokenVerifier:
    """Verifies ``Authorization: Bearer <jwt>`` tokens against the JWKS endpoint."""

    def __init__(self, config: ResolvedConfig) -> None:
        self._config = config
        self._jwks_client = PyJWKClient(
            config.jwks_uri,
            cache_keys=True,
            lifespan=config.jwks.cache_max_age or 300,
            max_cached_keys=config.jwks.max_cached_keys or 16,
            timeout=config.jwks.timeout_duration or 30,
        )

    def verify(
        self,
        token: str,
        *,
        audience: Optional[Audience] = None,
        issuer: Optional[str] = None,
    ) -> dict[str, Any]:
        """Resolve the signing key from the token's ``kid`` and verify it.

        Returns the decoded claims, or raises ``jwt.PyJWTError`` on failure.
        """
        signing_key = self._jwks_client.get_signing_key_from_jwt(token)
        options = {"verify_aud": True, "verify_iss": True}

        expected_audience = audience if audience is not None else self._config.audience
        expected_issuer = issuer if issuer is not None else self._config.issuer
        if expected_issuer is None:
            options["verify_iss"] = False

        return jwt.decode(
            token,
            signing_key.key,
            algorithms=["RS256", "ES256", "RS384", "RS512", "ES384", "ES512"],
            audience=expected_audience,
            issuer=expected_issuer,
            options=options,
        )


def extract_bearer_token(authorization_header: Optional[str]) -> Optional[str]:
    """Pull the token out of an ``Authorization: Bearer <token>`` header value."""
    if not authorization_header:
        return None
    parts = authorization_header.split(" ", 1)
    if len(parts) != 2:
        return None
    scheme, token = parts
    if scheme.lower() != "bearer":
        return None
    token = token.strip()
    return token or None


def extract_bearer_from_headers(headers: Mapping[str, str]) -> Optional[str]:
    """Convenience for header mappings (case-insensitive ``Authorization``)."""
    for key, value in headers.items():
        if key.lower() == "authorization":
            return extract_bearer_token(value)
    return None
