"""JWTAuth base class."""

from __future__ import annotations

from datetime import timedelta
from typing import Any, Iterable, cast

import jwt
from anyio import to_thread
from jwt.types import Options

from .constants import (
    DEFAULT_ALGORITHMS,
    DEFAULT_JWKS_LIFESPAN,
    DEFAULT_LEEWAY,
    DEFAULT_OPTIONS,
)
from .exceptions import JWTAuthError


class JWTAuth:
    """JWTAuth object for performing decode operations."""

    def __init__(
        self,
        jwks_url: str,
        jwks_lifespan: int = DEFAULT_JWKS_LIFESPAN,
        algorithms: list[str] | None = None,
        audience: str | Iterable[str] | None = None,
        issuer: str | Iterable[str] | None = None,
        leeway: int | float | timedelta = DEFAULT_LEEWAY,
        options: dict[str, Any] | None = None,
    ):
        """Create a JWTAuth object."""
        self.jwks_url = jwks_url
        self.jwks_lifespan = jwks_lifespan or DEFAULT_JWKS_LIFESPAN
        self.jwks_client = jwt.PyJWKClient(
            uri=self.jwks_url, lifespan=self.jwks_lifespan
        )

        # Decode parameters
        self.algorithms = algorithms or DEFAULT_ALGORITHMS
        self.audience = audience
        self.issuer = issuer
        self.leeway = leeway
        options_data = cast(
            Options, options if options is not None else DEFAULT_OPTIONS
        )
        self.options = Options(**options_data)

    def get_token(self, token_string: str) -> dict[str, Any]:
        """Get decoded token from string."""
        try:
            signing_key = self.jwks_client.get_signing_key_from_jwt(token_string)
        except jwt.PyJWKClientConnectionError as exc:
            raise JWTAuthError(
                "Failed to fetch from JWKS URL",
                code="jwks_connect_error",
            ) from exc
        except jwt.PyJWKClientError as exc:
            raise JWTAuthError(
                f"JWK Client Error {exc}", code="jwk_client_error"
            ) from exc
        except jwt.PyJWTError as exc:
            raise JWTAuthError(
                "Failed to get signing key from JWT.",
                code="invalid_token",
            ) from exc
        return self._decode_token(token_string, key=signing_key.key)

    async def aget_token(self, token_string: str) -> dict[str, Any]:
        """Asynchronously get decoded token from string."""
        try:
            signing_key = await to_thread.run_sync(
                self.jwks_client.get_signing_key_from_jwt, token_string
            )
        except jwt.PyJWKClientConnectionError as exc:
            raise JWTAuthError(
                "Failed to fetch from JWKS URL",
                code="jwks_connect_error",
            ) from exc
        except jwt.PyJWKClientError as exc:
            raise JWTAuthError(
                f"JWK Client Error {exc}", code="jwk_client_error"
            ) from exc
        except jwt.PyJWTError as exc:
            raise JWTAuthError(
                "Failed to get signing key from JWT.",
                code="invalid_token",
            ) from exc
        return self._decode_token(token_string, key=signing_key.key)

    def _decode_token(self, token_string: str, key: str = "") -> dict[str, Any]:
        """Helper to decode token from string."""
        try:
            payload = jwt.decode(
                token_string,
                key=key,
                algorithms=self.algorithms,
                options=self.options,
                audience=self.audience,
                issuer=None,
                leeway=self.leeway,
            )
            # Add custom issuer validation
            if "verify_iss" in self.options:
                self._validate_iss(payload)
            return payload
        except jwt.PyJWKClientError as exc:
            raise JWTAuthError(
                f"JWK Client Error {exc}", code="jwk_client_error"
            ) from exc
        except jwt.PyJWTError as exc:
            raise JWTAuthError(
                "The access token provided is expired, revoked, malformed, "
                "or invalid for other reasons.",
                code="invalid_token",
            ) from exc

    def _validate_iss(self, payload: dict[str, Any]) -> None:
        """Helper to validate token issuer."""
        if self.issuer is None:
            return

        try:
            iss = payload["iss"]
        except KeyError as exc:
            raise jwt.MissingRequiredClaimError(
                "iss"
            ) from exc  # mypy: ignore[no-untyped-call]

        issuer = self.issuer

        if isinstance(issuer, str):
            issuer = (issuer,)

        if iss not in issuer:
            raise jwt.InvalidIssuerError("Invalid issuer")