import hashlib
from authlib.common.encoding import to_bytes, to_unicode, urlsafe_b64encode
from ..rfc6749.errors import InvalidRequestError, InvalidGrantError


def create_s256_code_challenge(code_verifier):
    """Create S256 code_challenge with the given code_verifier."""
    data = hashlib.sha256(to_bytes(code_verifier, 'ascii')).digest()
    return to_unicode(urlsafe_b64encode(data))


def compare_plain_code_challenge(code_verifier, code_challenge):
    # If the "code_challenge_method" from Section 4.3 was "plain",
    # they are compared directly
    return code_verifier == code_challenge


def compare_s256_code_challenge(code_verifier, code_challenge):
    # BASE64URL-ENCODE(SHA256(ASCII(code_verifier))) == code_challenge
    return create_s256_code_challenge(code_verifier) == code_challenge


class CodeChallenge(object):
    """CodeChallenge extension to Authorization Code Grant. It is used to
    improve the security of Authorization Code flow for public clients by
    sending extra "code_challenge" and "code_verifier" to the authorization
    server.

    The AuthorizationCodeGrant SHOULD save the code_challenge and
    code_challenge_method into database when create_authorization_code. Then
    register this extension via::

        server.register_grant(
            AuthorizationCodeGrant,
            [CodeChallenge(required=True)]
        )
    """
    #: defaults to "plain" if not present in the request
    DEFAULT_CODE_CHALLENGE_METHOD = 'plain'
    #: supported ``code_challenge_method``
    SUPPORTED_CODE_CHALLENGE_METHOD = ['plain', 'S256']

    CODE_CHALLENGE_METHODS = {
        'plain': compare_plain_code_challenge,
        'S256': compare_s256_code_challenge,
    }

    def __init__(self, required=True):
        self.required = required

    def __call__(self, grant):
        grant.register_hook(
            'after_validate_authorization_request',
            self.validate_code_challenge,
        )
        grant.register_hook(
            'after_validate_token_request',
            self.validate_code_verifier,
        )

    def validate_code_challenge(self, grant):
        request = grant.request
        challenge = request.args.get('code_challenge')
        method = request.args.get('code_challenge_method')
        if not challenge and not method:
            return

        if not challenge:
            raise InvalidRequestError('Missing "code_challenge"')

        if method and method not in self.SUPPORTED_CODE_CHALLENGE_METHOD:
            raise InvalidRequestError('Unsupported "code_challenge_method"')

    def validate_code_verifier(self, grant):
        request = grant.request
        verifier = request.form.get('code_verifier')

        # public client MUST verify code challenge
        if self.required and request.auth_method == 'none' and not verifier:
            raise InvalidRequestError('Missing "code_verifier"')

        authorization_code = request.credential
        challenge = self.get_authorization_code_challenge(authorization_code)

        # ignore, it is the normal RFC6749 authorization_code request
        if not challenge:
            return

        # challenge exists, code_verifier is required
        if not verifier:
            raise InvalidRequestError('Missing "code_verifier"')

        # 4.6. Server Verifies code_verifier before Returning the Tokens
        method = self.get_authorization_code_challenge_method(authorization_code)
        if method is None:
            method = self.DEFAULT_CODE_CHALLENGE_METHOD

        func = self.CODE_CHALLENGE_METHODS.get(method)
        if not func:
            raise RuntimeError('No verify method for "{}"'.format(method))

        # If the values are not equal, an error response indicating
        # "invalid_grant" MUST be returned.
        if not func(verifier, challenge):
            raise InvalidGrantError(description='Code challenge failed.')

    def get_authorization_code_challenge(self, authorization_code):
        """Get "code_challenge" associated with this authorization code.
        Developers MAY re-implement it in subclass, the default logic::

            def get_authorization_code_challenge(self, authorization_code):
                return authorization_code.code_challenge

        :param authorization_code: the instance of authorization_code
        """
        return authorization_code.code_challenge

    def get_authorization_code_challenge_method(self, authorization_code):
        """Get "code_challenge_method" associated with this authorization code.
        Developers MAY re-implement it in subclass, the default logic::

            def get_authorization_code_challenge_method(self, authorization_code):
                return authorization_code.code_challenge_method

        :param authorization_code: the instance of authorization_code
        """
        return authorization_code.code_challenge_method