import json
from functools import wraps
from urllib.request import urlopen

from flask import _request_ctx_stack, request
from jose import JWTError, jwt

from slz_api_service.const import AUTH0_ALGORITHMS, AUTH0_API_AUDIENCE, AUTH0_DOMAIN
from slz_api_service.errors import AuthError, Codes


class ProxyResourceProtector:
    """Created in the spirit of AuthLib's ResourceProtector_ class for Flask but to be used as a
    proxy to just protect routes by validating the auth token provided in the request.

    Calling this returns a decorator that can be used to protect routes with an optional scope
    specified.

    .. _ResourceProtector:
       https://github.com/lepture/authlib/blob/master/authlib/flask/oauth2/resource_protector.py#L19

    Example:

    .. code-block:: python

        requires_oauth = ResourceProtector()

        @app.route('/user')
        @require_oauth('profile')
        def user_profile():
            user = User.query.get(current_token.user_id)
            return jsonify(user.to_dict())

    """

    @classmethod
    def get_token_auth_header(cls) -> str:
        """
        Returns:
            token string part from the Authorization Header
        """
        auth_header = request.headers.get('Authorization', None)
        if not auth_header:
            raise AuthError({
                'code': Codes.authorization_header_missing.value,
                'description': 'Authorization header is required to access this resource',
            })

        parts = auth_header.split()

        if parts[0].lower() != 'bearer':
            raise AuthError({
                'code': Codes.invalid_header.value,
                'description': 'Authorization header must begin with the token type: Bearer',
            })
        elif len(parts) == 1:
            raise AuthError({
                'code': Codes.invalid_header.value,
                'description': 'Token not found in Authorization header',
            })
        elif len(parts) > 2:
            raise AuthError({
                'code': Codes.invalid_header.value,
                'description': 'Authorization header must only include Bearer and token',
            })

        token = parts[1]
        return token

    @classmethod
    def get_rsa_key(cls, token: str) -> dict:
        jsonurl = urlopen('https://' + AUTH0_DOMAIN + '/.well-known/jwks.json')
        jwks = json.loads(jsonurl.read())
        try:
            unverified_header = jwt.get_unverified_header(token)
        except JWTError as e:
            raise AuthError({
                'code': Codes.invalid_token.value,
                'description': str(e),
            })
        rsa_key = {}
        for key in jwks['keys']:
            if key['kid'] == unverified_header['kid']:
                rsa_key = {
                    'kty': key['kty'],
                    'kid': key['kid'],
                    'use': key['use'],
                    'n': key['n'],
                    'e': key['e'],
                }
        return rsa_key

    @classmethod
    def decode_token(cls, token: str, rsa_key: dict = None):
        if not rsa_key:
            rsa_key = cls.get_rsa_key(token)
        try:
            payload = jwt.decode(
                token,
                rsa_key,
                algorithms=AUTH0_ALGORITHMS,
                audience=AUTH0_API_AUDIENCE,
                issuer='https://' + AUTH0_DOMAIN + '/',
            )
        except jwt.ExpiredSignatureError:
            raise AuthError({
                'code': Codes.token_expired.value,
                'description': 'Auth token is expired',
            })
        except jwt.JWTClaimsError:
            raise AuthError({
                'code': Codes.invalid_claims.value,
                'description': 'Incorrect claims, please check the audience and issuer',
            })
        except Exception:
            raise AuthError({
                'code': Codes.invalid_header.value,
                'description': 'Unable to parse authentication token',
            })
        return payload

    def __call__(self, scope=None):
        def wrapper(f):
            """Decorator to be used on routes requiring authentication and authorization.
            Determines if the Access Token is valid. See this class documentation for example usage.

            Raises:
                AuthError: If the token is invalid
            """

            @wraps(f)
            def decorated(*args, **kwargs):
                token = self.get_token_auth_header()
                rsa_key = self.get_rsa_key(token)
                if rsa_key:
                    payload = self.decode_token(token, rsa_key)

                    if scope and not self.requires_scope(scope):
                        raise AuthError({
                            'code': Codes.invalid_scope.value,
                            'description': f'Client is missing required scope for this resource: '
                                           f'"{scope}"',
                        })

                    _request_ctx_stack.top.current_user = payload
                    return f(*args, **kwargs)
                raise AuthError({
                    'code': Codes.invalid_header.value,
                    'description': 'Unable to find appropriate key',
                }, 401)

            return decorated

        return wrapper

    @classmethod
    def requires_scope(cls, required_scope: str) -> bool:
        """Determines if the required scope is present in the Access Token

        Args:
            required_scope: The scope required to access the resource

        Returns:
            True if the authenticated client has access to the provided ``required_scope``
        """
        token = cls.get_token_auth_header()
        unverified_claims = jwt.get_unverified_claims(token)
        if unverified_claims.get('scope'):
            token_scopes = unverified_claims['scope'].split()
            for token_scope in token_scopes:
                if token_scope == required_scope:
                    return True
        return False