from __future__ import annotations
from datetime import datetime

import hashlib
from typing import Callable, Optional

from psion.oauth2.exceptions import InvalidGrant, InvalidRequest
from psion.oauth2.mixins import AuthorizationCodeMixin, ClientMixin, UserMixin
from psion.webtools import (
    base64url_encode,
    secret_token,
    to_bytes,
    to_string,
    urlencode,
)

from .base import BaseGrant


def plain_challenge(challenge: str, verifier: str):
    return challenge == verifier


def S256_challenge(challenge: str, verifier: str):
    hashed_verifier = to_string(
        base64url_encode(hashlib.sha256(to_bytes(verifier, "ascii")).digest())
    )
    return challenge == hashed_verifier


class AuthorizationCodeGrant(BaseGrant):
    """
    Implementation of the Authorization Code Grant described in the
    `OAuth 2.1 Spec <https://tools.ietf.org/html/draft-parecki-oauth-v2-1#section-4.1>`_.

    In this grant, the client **MUST** obtain an authorization grant from the RO
    and exchange it for an access token. This implementation uses PKCE by default,
    and enforces its use every time.

    This grant allows all types of clients to get grants and tokens.
    Depending on the client, it **MAY** issue refresh tokens as well.

    The PKCE methods supported are `plain` and `S256` and, even though `plain`
    is supported, its usage is **DISCOURAGED**, since it does not provide
    enough security for the flow.
    """

    __response_type__: str = "code"
    __grant_type__: str = "authorization_code"

    _challenges: dict[str, Callable[[str, str], bool]] = {
        "plain": plain_challenge,
        "S256": S256_challenge,
    }

    async def authorize(self, data: dict, client: ClientMixin, user: UserMixin) -> str:
        """
        Validates the `Authorization Request` and construes the `Authorization Response`
        containing the "code" and the "state", if the latter was provided by the Client.

        From the specification at
        `<https://tools.ietf.org/html/draft-parecki-oauth-v2-1-03#section-4.1.2>`_::

            "code": REQUIRED. The authorization code generated by the
                authorization server. The authorization code MUST expire shortly
                after it is issued to mitigate the risk of leaks. A maximum
                authorization code lifetime of 10 minutes is RECOMMENDED. The
                client MUST NOT use the authorization code more than once. If an
                authorization code is used more than once, the authorization
                server MUST deny the request and SHOULD revoke (when possible) all
                tokens previously issued based on that authorization code. The
                authorization code is bound to the client identifier and redirect
                URI.

            "state": REQUIRED if the "state" parameter was present in the client
                authorization request. The exact value received from the client.

            For example, the authorization server redirects the user-agent by
            sending the following HTTP response:

            HTTP/1.1 302 Found
            Location: https://client.example.com/cb?code=SplxlOBeZQQYbYS6WxSbIA&state=xyz

        :param data: Data of the Authorization Request.
        :type data: dict

        :param client: Client requesting authorization.
        :type client: ClientMixin

        :param user: Currently authenticated User.
        :type user: UserMixin

        :return: URL Encoded Authorization Response data.
        :rtype: str
        """

        data = await self._validate_authorization_request(data)
        self._validate_requested_scopes(data["scopes"], client, data.get("state"))

        code = secret_token(64)
        await self.adapter.save_authorization_code(code, data, client, user)

        response = {"code": code, "state": data.get("state")}

        async for hook in self._execute_hook(
            "authorization_response", code, data, client, user
        ):
            if hook:
                response.update(hook)

        return self._create_authorization_response(data["redirect_uri"], response)

    async def _validate_authorization_request(self, data: dict) -> dict:
        """
        Validates the incoming data from the `Client` to ensure
        that **ALL** the required parameters were provided.

        From the specification at
        `<https://tools.ietf.org/html/draft-parecki-oauth-v2-1-03#section-4.1.1.3>`_,
        (fields marked in italic are slightly modified
        to fit the Framework's requirements)::

            "response_type": REQUIRED. Value MUST be set to "code".

            "client_id": REQUIRED. The client identifier as described in Section 2.2.

            "code_challenge": REQUIRED. Code challenge.

            "code_challenge_method": *REQUIRED*. Code verifier transformation
                method is "S256" or "plain".

            "redirect_uri": *REQUIRED*. As described in Section 3.1.2.

            "scope": *REQUIRED*. The scope of the access request
                as described by Section 3.3.

            "state": OPTIONAL. An opaque value used by the client to maintain
                state between the request and callback. The authorization server
                includes this value when redirecting the user-agent back to the
                client.

            GET /authorize?response_type=code&client_id=s6BhdRkqt3&state=xyz
                &redirect_uri=https%3A%2F%2Fclient%2Eexample%2Ecom%2Fcb
                &code_challenge=6fdkQaPm51l13DSukcAH3Mdx7_ntecHYd1vi3n0hMZY
                &code_challenge_method=S256 HTTP/1.1
            Host: server.example.com

        :param data: Data of the Authorization Request.
        :type data: dict

        :return: Validated and reformatted Authorization Request data.
        :rtype: dict
        """

        state: Optional[str] = data.get("state")
        code_challenge: str = data.get("code_challenge")
        code_challenge_method: str = data.get("code_challenge_method")
        scope: str = data.get("scope")

        if state:
            if not state or not isinstance(state, str):
                raise InvalidRequest(description='Invalid parameter "state".')

        if not code_challenge or not isinstance(code_challenge, str):
            raise InvalidRequest(description="Invalid code challenge.", state=state)

        if len(code_challenge) < 43 or len(code_challenge) > 128:
            raise InvalidRequest(
                description="The length of the code challenge MUST be between [43, 128] bytes long.",
                state=state,
            )

        if code_challenge_method not in self._challenges.keys():
            raise InvalidRequest(
                description=f'Unknown code challenge "{code_challenge_method}".',
                state=state,
            )

        if not scope or not isinstance(scope, str):
            raise InvalidRequest(description='Invalid parameter "scope".', state=state)

        response = {
            "code_challenge": code_challenge,
            "code_challenge_method": code_challenge_method,
            "redirect_uri": data["redirect_uri"],
            "scopes": scope.split(),
            "state": state,
        }

        async for hook in self._execute_hook("authorization_request", data):
            if hook:
                response.update(hook)

        return response

    async def token(self, data: dict, client: ClientMixin) -> dict:
        """
        Validates the provided `Authorization Code` to check if its data matches
        both the `Client` and the `User`, then issues a new `Access Token` and,
        if allowed, a new `Refresh Token`, with both bound to the `Client` and `User`.

        Whether it succeeds or fails to issue an `Access Token`, the current
        `Authorization Code` **WILL** be deleted at the end of the flow.
        This prevents both `Replay Attacks` and the generation of multiple
        `Access Tokens` against the `Authorization Code`.

        From the specification at
        `<https://tools.ietf.org/html/draft-parecki-oauth-v2-1-03#section-5.1>`_::

            The authorization server issues an access token and optional refresh
            token, and constructs the response by adding the following parameters
            to the entity-body of the HTTP response with a 200 (OK) status code:

            "access_token": REQUIRED. The access token issued by the
                authorization server.

            "token_type": REQUIRED. The type of the token issued as described
                in Section 7.1. Value is case insensitive.

            "expires_in": RECOMMENDED. The lifetime in seconds of the access
                token. For example, the value "3600" denotes that the access
                token will expire in one hour from the time the response was
                generated. If omitted, the authorization server SHOULD provide
                the expiration time via other means or document the default value.

            "refresh_token": OPTIONAL. The refresh token, which can be used to
                obtain new access tokens using the same authorization grant as
                described in Section 6.

            "scope": OPTIONAL, if identical to the scope requested by the
                client; otherwise, REQUIRED. The scope of the access token as
                described by Section 3.3.

            The parameters are included in the entity-body of the HTTP response
            using the "application/json" media type as defined by [RFC7159]. The
            parameters are serialized into a JavaScript Object Notation (JSON)
            structure by adding each parameter at the highest structure level.
            Parameter names and string values are included as JSON strings.
            Numerical values are included as JSON numbers. The order of
            parameters does not matter and can vary.

            The authorization server MUST include the HTTP "Cache-Control"
            response header field [RFC7234] with a value of "no-store" in any
            response containing tokens, credentials, or other sensitive
            information, as well as the "Pragma" response header field [RFC7234]
            with a value of "no-cache".

            For example:

            HTTP/1.1 200 OK
            Content-Type: application/json
            Cache-Control: no-store
            Pragma: no-cache

            {
                "access_token":"2YotnFZFEjr1zCsicMWpAA",
                "token_type":"Bearer",
                "expires_in":3600,
                "refresh_token":"tGzv3JOkF0XG5Qx2TlKWIA",
                "example_parameter":"example_value"
            }

            The client MUST ignore unrecognized value names in the response.  The
            sizes of tokens and other values received from the authorization
            server are left undefined. The client should avoid making
            assumptions about value sizes. The authorization server SHOULD
            document the size of any value it issues.

        :param data: Data of the Token Request.
        :type data: dict

        :param client: Client requesting a new Token.
        :type client: ClientMixin

        :return: Access Token and its metadata, optionally with a Refresh Token.
        :rtype: dict
        """

        try:
            data = await self._validate_token_request(data)
            code = await self.adapter.get_authorization_code(data["code"])

            if not code:
                raise InvalidRequest(description="Invalid Authorization Grant.")

            self._validate_authorization_code(data, code, client)

            user = await self.adapter.find_user(code.get_user_id())

            if not user:
                raise InvalidRequest(description="No user found for this code.")

            access_token = await self.adapter.create_access_token(
                client, user, code.get_scopes()
            )

            refresh_token = (
                await self.adapter.create_refresh_token(client, user, code.get_scopes())
                if "refresh_token" in client.get_grant_types()
                else None
            )

            token = self._create_token(
                access_token,
                self.config.token_lifespan,
                refresh_token,
                code.get_scopes(),
            )

            async for hook in self._execute_hook(
                "token_response", token, data, client, user
            ):
                if hook:
                    token.update(hook)

            return token
        finally:
            # Prevents replay attacks by deleting the code once it is used,
            # whether it succeeds or fails.
            await self.adapter.delete_authorization_code(data["code"])

    async def _validate_token_request(self, data: dict) -> dict:
        """
        Validates the incoming data from the `Client` to ensure
        that **ALL** the required parameters were provided.

        From the specification at
        `<https://tools.ietf.org/html/draft-parecki-oauth-v2-1-03#section-4.1.3>`_::

            The client makes a request to the token endpoint by sending the
            following parameters using the "application/x-www-form-urlencoded"
            format per Appendix B with a character encoding of UTF-8 in the HTTP
            request entity-body:

            "grant_type": REQUIRED. Value MUST be set to "authorization_code".

            "code": REQUIRED. The authorization code received from the
                authorization server.

            "redirect_uri": REQUIRED, if the "redirect_uri" parameter was
                included in the authorization request as described in
                Section 4.1.1, and their values MUST be identical.

            "client_id": REQUIRED, if the client is not authenticating with the
                authorization server as described in Section 3.2.1.

            "code_verifier": REQUIRED, if the "code_challenge" parameter was
                included in the authorization request. MUST NOT be used
                otherwise. The original code verifier string.

            Confidential or credentialed clients MUST authenticate with the
            authorization server as described in Section 3.2.1.

            For example, the client makes the following HTTP request using TLS
            (with extra line breaks for display purposes only):

            POST /token HTTP/1.1
            Host: server.example.com
            Authorization: Basic czZCaGRSa3F0MzpnWDFmQmF0M2JW
            Content-Type: application/x-www-form-urlencoded

            grant_type=authorization_code&code=SplxlOBeZQQYbYS6WxSbIA
            &redirect_uri=https%3A%2F%2Fclient%2Eexample%2Ecom%2Fcb
            &code_verifier=3641a2d12d66101249cdf7a79c000c1f8c05d2aafcf14bf146497bed

            The authorization server MUST:

            * require client authentication for confidential and credentialed
                clients (or clients with other authentication requirements),

            * authenticate the client if client authentication is included,

            * ensure that the authorization code was issued to the authenticated
                confidential or credentialed client, or if the client is public,
                ensure that the code was issued to "client_id" in the request,

            * verify that the authorization code is valid,

            * verify that the "code_verifier" parameter is present if and only
                if a "code_challenge" parameter was present in the authorization
                request,

            * if a "code_verifier" is present, verify the "code_verifier" by
                calculating the code challenge from the received "code_verifier"
                and comparing it with the previously associated "code_challenge",
                after first transforming it according to the
                "code_challenge_method" method specified by the client, and

            * ensure that the "redirect_uri" parameter is present if the
                "redirect_uri" parameter was included in the initial authorization
                request as described in Section 4.1.1.3, and if included ensure
                that their values are identical.

        .. note:: The `client_id` is ignored in this flow, since it is already
            used to authenticate the `Client` if needed.

        :param data: Data of the Token Request.
        :type data: dict

        :return: Validated and reformatted Token Request data.
        :rtype: dict
        """

        if not data.get("code") or not isinstance(data.get("code"), str):
            raise InvalidRequest(description='Invalid parameter "code".')

        if not data.get("redirect_uri") or not isinstance(
            data.get("redirect_uri"), str
        ):
            raise InvalidRequest(description='Invalid parameter "redirect_uri".')

        if not data.get("code_verifier") or not isinstance(
            data.get("code_verifier"), str
        ):
            raise InvalidRequest(description='Invalid parameter "code_verifier".')

        response = {
            "code": data["code"],
            "redirect_uri": data["redirect_uri"],
            "code_verifier": data["code_verifier"],
        }

        async for hook in self._execute_hook("token_request", data):
            if hook:
                response.update(hook)

        return response

    def _validate_authorization_code(
        self, data: dict, code: AuthorizationCodeMixin, client: ClientMixin
    ) -> None:
        """
        Validates the data of the `Provided Authorization Code` against both
        the current `Client` and the `Stored Authorization Code`.

        :param data: Provided Authorization Code.
        :type data: AuthorizationCodeModel

        :param code: Stored Authorization Code.
        :type code: AuthorizationCodeMixin

        :param client: Current Client requesting an Access Token.
        :type client: ClientMixin

        :raises InvalidGrant: The data of the Client, the Request
            and the stored Authorization Code do not match.
        :raises InvalidRequest: The transformation method is not supported.
        """

        if data["redirect_uri"] not in client.get_redirect_uris():
            raise InvalidGrant(description="Invalid Redirect URI.")

        if client.get_client_id() != code.get_client_id():
            raise InvalidGrant(description="Mismatching Client ID.")

        if data["redirect_uri"] != code.get_redirect_uri():
            raise InvalidGrant(description="Mismatching Redirect URI.")

        method = self._challenges.get(code.get_code_challenge_method())

        if not method:
            raise InvalidRequest(description=f'Unknown transform "{method}".')

        if not method(code.get_code_challenge(), data["code_verifier"]):
            raise InvalidGrant(description="PKCE challenge failure.")

        if datetime.utcnow() >= code.get_expiration():
            raise InvalidGrant(description="Expired Authorization Code.")

    def _create_authorization_response(self, url: str, data: dict) -> str:
        """
        Creates the Redirect URI with the parameters of the Authorization Response
        to be returned back to the Client.

        :param url: Callback URL of the Client.
        :type url: str

        :param data: Data of the Authorization Response.
        :type data: dict

        :return: URL encoded data of the authorization response.
        :rtype: str
        """

        return urlencode(url, **data)
