"""JWT validation using JWKS."""

import jwt
from jwt.algorithms import RSAAlgorithm

from ..exceptions import AuthenticationError
from .claims import TokenClaims
from .jwks import JWKSClient


class JWTValidator:
    """Validates JWTs using JWKS from vikdata.

    This validator fetches public keys from a JWKS endpoint and uses them
    to verify JWT signatures. It supports RS256 algorithm and handles
    key rotation automatically by refreshing JWKS when a key is not found.

    Attributes:
        jwks_client: JWKSClient instance for fetching keys
        audience: Expected audience claim value(s)
        issuer: Expected issuer claim value
        algorithms: Allowed JWT algorithms
    """

    def __init__(
        self,
        jwks_url: str,
        audience: str | list[str] = "vikdata",
        issuer: str = "https://api.vikdata.com",
        algorithms: list[str] | None = None,
        cache_ttl: int = 3600,
    ):
        """Initialize JWT validator.

        Args:
            jwks_url: URL to fetch JWKS from
            audience: Expected audience (single value or list)
            issuer: Expected issuer
            algorithms: Allowed algorithms (default: ["RS256"])
            cache_ttl: JWKS cache TTL in seconds
        """
        self.jwks_client = JWKSClient(jwks_url, cache_ttl=cache_ttl)
        self.audience = audience
        self.issuer = issuer
        self.algorithms = algorithms or ["RS256"]

    async def validate(self, token: str) -> TokenClaims:
        """Validate JWT and return claims.

        This method:
        1. Decodes the token header to get the key ID (kid)
        2. Fetches the JWKS and finds the matching public key
        3. Verifies the signature using the public key
        4. Validates standard claims (exp, aud, iss)
        5. Returns parsed claims

        Args:
            token: JWT string to validate

        Returns:
            TokenClaims object with validated claims

        Raises:
            AuthenticationError: If token is invalid, expired, or has wrong audience/issuer
        """
        try:
            # Get unverified header to find kid
            header = jwt.get_unverified_header(token)
            kid = header.get("kid")

            if not kid:
                raise AuthenticationError("Token missing 'kid' in header")

            # Get JWKS and find key
            jwks = await self.jwks_client.get_jwks()
            try:
                key_data = self.jwks_client.get_key(jwks, kid)
            except KeyError:
                # Key not found, try refreshing JWKS (key rotation scenario)
                jwks = await self.jwks_client.get_jwks(force_refresh=True)
                try:
                    key_data = self.jwks_client.get_key(jwks, kid)
                except KeyError:
                    raise AuthenticationError(f"Unknown key ID: {kid}")

            # Convert JWK to public key
            public_key = RSAAlgorithm.from_jwk(key_data)

            # Verify and decode
            payload = jwt.decode(
                token,
                public_key,
                algorithms=self.algorithms,
                audience=self.audience,
                issuer=self.issuer,
            )

            return TokenClaims.from_payload(payload)

        except jwt.ExpiredSignatureError:
            raise AuthenticationError("Token has expired")
        except jwt.InvalidAudienceError:
            raise AuthenticationError("Invalid token audience")
        except jwt.InvalidIssuerError:
            raise AuthenticationError("Invalid token issuer")
        except jwt.DecodeError as e:
            raise AuthenticationError(f"Invalid token format: {e}")
        except jwt.InvalidTokenError as e:
            raise AuthenticationError(f"Invalid token: {e}")

    async def validate_without_expiry_check(self, token: str) -> TokenClaims:
        """Validate JWT without checking expiration.

        Useful for token refresh flows where you need to read expired tokens.

        Args:
            token: JWT string to validate

        Returns:
            TokenClaims object with validated claims (may be expired)

        Raises:
            AuthenticationError: If token is invalid or has wrong audience/issuer
        """
        try:
            header = jwt.get_unverified_header(token)
            kid = header.get("kid")

            if not kid:
                raise AuthenticationError("Token missing 'kid' in header")

            jwks = await self.jwks_client.get_jwks()
            try:
                key_data = self.jwks_client.get_key(jwks, kid)
            except KeyError:
                jwks = await self.jwks_client.get_jwks(force_refresh=True)
                try:
                    key_data = self.jwks_client.get_key(jwks, kid)
                except KeyError:
                    raise AuthenticationError(f"Unknown key ID: {kid}")

            public_key = RSAAlgorithm.from_jwk(key_data)

            # Decode with expiry verification disabled
            payload = jwt.decode(
                token,
                public_key,
                algorithms=self.algorithms,
                audience=self.audience,
                issuer=self.issuer,
                options={"verify_exp": False},
            )

            return TokenClaims.from_payload(payload)

        except jwt.InvalidAudienceError:
            raise AuthenticationError("Invalid token audience")
        except jwt.InvalidIssuerError:
            raise AuthenticationError("Invalid token issuer")
        except jwt.DecodeError as e:
            raise AuthenticationError(f"Invalid token format: {e}")
        except jwt.InvalidTokenError as e:
            raise AuthenticationError(f"Invalid token: {e}")
