""" JWT decoder utility functions. """
import logging
import sys

from jwkest.jwk import KEYS
from jwkest.jws import JWS
import jwt
from semantic_version import Version
from django.conf import settings
from rest_framework_jwt.settings import api_settings

from edx_rest_framework_extensions.settings import get_jwt_issuers, get_first_jwt_issuer

logger = logging.getLogger(__name__)


class JwtTokenVersion(object):
    latest_supported = '1.1.0'

    starting_version = '1.0.0'
    added_version = '1.1.0'
    added_is_restricted = '1.1.0'
    added_filters = '1.1.0'


def jwt_decode_handler(token):
    """
    Decodes a JSON Web Token (JWT).

    Notes:
        * Requires "exp" and "iat" claims to be present in the token's payload.
        * Aids debugging by logging InvalidTokenError log entries when decoding fails.

    Examples:
        Use with `djangorestframework-jwt <https://getblimp.github.io/django-rest-framework-jwt/>`_, by changing
        your Django settings:

        .. code-block:: python

            JWT_AUTH = {
                'JWT_DECODE_HANDLER': 'edx_rest_framework_extensions.jwt_decoder.jwt_decode_handler',
                'JWT_ISSUER': 'https://the.jwt.issuer',
                'JWT_SECRET_KEY': 'the-jwt-secret-key',  (defaults to settings.SECRET_KEY)
                'JWT_AUDIENCE': 'the-jwt-audience',
                'JWT_PUBLIC_SIGNING_JWK_SET': 'the-jwk-set-of-public-signing-keys',
            }

    Args:
        token (str): JWT to be decoded.

    Returns:
        dict: Decoded JWT payload.

    Raises:
        MissingRequiredClaimError: Either the exp or iat claims is missing from the JWT payload.
        InvalidTokenError: Decoding fails.
    """
    jwt_issuer = get_first_jwt_issuer()
    _verify_jwt_signature(token, jwt_issuer)
    decoded_token = _decode_and_verify_token(token, jwt_issuer)
    return _set_token_defaults(decoded_token)


def decode_jwt_scopes(token):
    """
    Decode the JWT and return the scopes claim.
    """
    return jwt_decode_handler(token).get('scopes', [])


def decode_jwt_is_restricted(token):
    """
    Decode the JWT and return the is_restricted claim.
    """
    return jwt_decode_handler(token)['is_restricted']


def decode_jwt_filters(token):
    """
    Decode the JWT, parse the filters claim, and return a
    list of (filter_type, filter_value) tuples.
    """
    return [jwt_filter.split(':') for jwt_filter in jwt_decode_handler(token)['filters']]


def _set_token_defaults(token):
    """
    Returns an updated token that includes default values for
    fields that were introduced since the token was created
    by checking its version number.
    """
    def _verify_version(jwt_version):
        supported_version = Version(
            settings.JWT_AUTH.get('JWT_SUPPORTED_VERSION', JwtTokenVersion.latest_supported)
        )
        if jwt_version.major > supported_version.major:
            logger.info('Token decode failed due to unsupported JWT version number [%s]', str(jwt_version))
            raise jwt.InvalidTokenError

    def _get_and_set_version(token):
        """
        Tokens didn't always contain a version number so we
        default to a nominal starting number.
        """
        if 'version' not in token:
            token['version'] = str(JwtTokenVersion.starting_version)

        return Version(token['version'])

    def _set_is_restricted(token, token_version):
        """
        We can safely default to False since all "restricted" tokens
        created prior to this version were always created as expired
        tokens. Expired tokens would not validate and so would
        not get as this far into the decoding process.
        """
        if token_version < Version(JwtTokenVersion.added_is_restricted):
            token['is_restricted'] = False

    def _set_filters(token, token_version):
        """
        We can safely default to an empty list of filters since
        previously created tokens were either "restricted" (always
        expired) or had full access.
        """
        if token_version < Version(JwtTokenVersion.added_filters):
            token['filters'] = []

    token_version = _get_and_set_version(token)
    _verify_version(token_version)
    _set_is_restricted(token, token_version)
    _set_filters(token, token_version)
    return token


def _verify_jwt_signature(token, jwt_issuer):
    key_set = _get_signing_jwk_key_set(jwt_issuer)

    try:
        _ = JWS().verify_compact(token, key_set)
    except Exception as exc:
        logger.exception('Token verification failed.')
        exc_info = sys.exc_info()
        raise jwt.InvalidTokenError(exc_info[2])


def _decode_and_verify_token(token, jwt_issuer):
    options = {
        'require_exp': True,
        'require_iat': True,

        'verify_exp': api_settings.JWT_VERIFY_EXPIRATION,
        'verify_aud': settings.JWT_AUTH.get('JWT_VERIFY_AUDIENCE', True),
        'verify_iss': False,  # TODO (ARCH-204): manually verify until issuer is configured correctly.
        'verify_signature': False,  # Verified with JWS already
    }

    decoded_token = jwt.decode(
        token,
        jwt_issuer['SECRET_KEY'],
        api_settings.JWT_VERIFY,
        options=options,
        leeway=api_settings.JWT_LEEWAY,
        audience=jwt_issuer['AUDIENCE'],
        issuer=jwt_issuer['ISSUER'],
        algorithms=[api_settings.JWT_ALGORITHM],
    )

    # TODO (ARCH-204): verify issuer manually until it is properly configured.
    token_issuer = decoded_token.get('iss')
    issuer_matched = any(issuer['ISSUER'] == token_issuer for issuer in get_jwt_issuers())
    if not issuer_matched:
        logger.info('Token decode failed due to mismatched issuer [%s]', token_issuer)
        raise jwt.InvalidTokenError('%s is not a valid issuer.', token_issuer)

    return decoded_token


def _get_signing_jwk_key_set(jwt_issuer):
    """
    Returns a JWK Keyset containing all active keys that are configured
    for verifying signatures.
    """
    key_set = KEYS()

    # asymmetric keys
    signing_jwk_set = settings.JWT_AUTH.get('JWT_PUBLIC_SIGNING_JWK_SET')
    if signing_jwk_set:
        key_set.load_jwks(signing_jwk_set)

    # symmetric key
    key_set.add({'key': jwt_issuer['SECRET_KEY'], 'kty': 'oct'})

    return key_set
