from collections import OrderedDict
from functools import lru_cache
from typing import Any, Dict, Iterable, Optional, Sequence, Union

import jwt
import jwt.api_jwt
from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicKey
from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PublicKey
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
from jwt import PyJWK
from jwt.exceptions import InvalidAlgorithmError
from jwt.types import Options

from atlassian_jwt_auth import KeyIdentifier, algorithms, exceptions, key
from atlassian_jwt_auth.key import BasePublicKeyRetriever

AllowedPublicKeys = Union[
    RSAPublicKey, EllipticCurvePublicKey, Ed25519PublicKey, Ed448PublicKey
]


@lru_cache(maxsize=10)
def _load_public_key(
    algorithms: Sequence[str], public_key: str, algorithm: Optional[str]
) -> Any:
    """Returns a public key object instance given the public key and
    algorithm.

    This has been extracted out of JWTAuthVerifier to avoid possible memory
    leaks via retained instance references.
    """
    if isinstance(public_key, (RSAPublicKey, EllipticCurvePublicKey)):
        return public_key
    if algorithm is None or algorithm not in algorithms:
        raise InvalidAlgorithmError("The specified alg value is not allowed")
    py_jws = jwt.api_jws.PyJWS(algorithms=algorithms)
    alg_obj = py_jws._algorithms[algorithm]
    return alg_obj.prepare_key(public_key)


class JWTAuthVerifier(object):
    """This class can be used to verify a JWT."""

    def __init__(
        self, public_key_retriever: BasePublicKeyRetriever, **kwargs: Any
    ) -> None:
        self.public_key_retriever = public_key_retriever
        self.algorithms = algorithms.get_permitted_algorithm_names()
        self._seen_jti: OrderedDict[str, None] = OrderedDict()
        self._subject_should_match_issuer = kwargs.get(
            "subject_should_match_issuer", True
        )
        self._check_jti_uniqueness = kwargs.get("check_jti_uniqueness", False)

    def verify_jwt(
        self,
        a_jwt: Union[str, bytes],
        audience: str,
        leeway: int = 0,
        **requests_kwargs: Any,
    ) -> Dict[Any, Any]:
        """Verify if the token is correct

        Returns:
             dict: the claims of the given jwt if verification is successful.

        Raises:
            ValueError: if verification failed.
        """
        key_identifier = key._get_key_id_from_jwt_header(a_jwt)
        public_key = self._retrieve_pub_key(key_identifier, requests_kwargs)

        alg = jwt.get_unverified_header(a_jwt).get("alg", None)
        public_key_obj = self._load_public_key(public_key, alg or "RS256")
        return self._decode_jwt(
            a_jwt, key_identifier, public_key_obj, audience=audience, leeway=leeway
        )

    def _retrieve_pub_key(
        self, key_identifier: Union[KeyIdentifier, str], requests_kwargs: Any
    ) -> str:
        return self.public_key_retriever.retrieve(key_identifier, **requests_kwargs)

    def _load_public_key(self, public_key: str, algorithm: Optional[str]) -> Any:
        """Returns a public key object instance given the public key and
        algorithm.
        """
        return _load_public_key(tuple(self.algorithms), public_key, algorithm)

    def _decode_jwt(
        self,
        a_jwt: Union[str, bytes],
        key_identifier: KeyIdentifier,
        jwt_key: Union[AllowedPublicKeys, PyJWK, str, bytes],
        audience: Optional[Union[str, Iterable[str]]] = None,
        leeway: int = 0,
    ) -> Dict[Any, Any]:
        """Decode JWT and check if it's valid"""
        options: Options = Options(
            verify_signature=True,
            require=["exp", "iat"],
            verify_exp=True,
            verify_iat=True,
        )

        claims = jwt.decode(
            a_jwt,
            key=jwt_key,
            algorithms=self.algorithms,
            options=options,
            audience=audience,
            leeway=leeway,
        )

        if (
            not key_identifier.key_id.startswith("%s/" % claims["iss"])
            and key_identifier.key_id != claims["iss"]
        ):
            raise ValueError("Issuer does not own the supplied public key")

        if self._subject_should_match_issuer and (
            claims.get("sub") and claims["iss"] != claims["sub"]
        ):
            raise exceptions.SubjectDoesNotMatchIssuerException(
                "Issuer does not match the subject."
            )

        _aud = claims.get("aud", None)
        if _aud is None:
            _msg = (
                "Claims validity, the aud claim must be provided and cannot be empty."
            )
            raise KeyError(_msg)
        _exp = int(claims["exp"])
        _iat = int(claims["iat"])
        if _exp - _iat > 3600:
            _msg = "Claims validity, '%s', exceeds the maximum 1 hour." % (_exp - _iat)
            raise ValueError(_msg)
        _jti = claims["jti"]
        if self._check_jti_uniqueness:
            self._check_jti(_jti)
        return claims

    def _check_jti(self, jti: str) -> None:
        """Checks that the given jti has not been already been used."""
        if jti in self._seen_jti:
            raise exceptions.JtiUniquenessException(
                "The jti, '%s', has already been used." % jti
            )
        self._seen_jti[jti] = None
        while len(self._seen_jti) > 1000:
            self._seen_jti.popitem(last=False)
