# -*- coding: utf-8 -*-
#
# Copyright 2017 Gehirn Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import hashlib
from typing import (
    Any,
    Callable,
    Dict,
    List,
)

from ghoauth.oauth2.errors import (
    FatalOAuth2Error,
    InvalidGrantError,
    InvalidRequestError,
    InvalidScopeError,
    ServerError,
)
from ghoauth.oauth2.grant_types.abc import (
    AbstractBaseGrantTypeHandler,
    AbstractBaseResponseTypeHandler,
)
from ghoauth.oauth2.request import OAuth2RequestMixin
from ghoauth.oauth2.token import (
    authorization_code_factory,
    bearer_token_factory,
)
from ghoauth.uriutils import extend_uri_with_params
from jwt import JWT
from jwt.utils import b64encode

from .repository import AbstractOpenIDRepository
from .request import OpenIDConnectRequestMixin
from .validator import AbstractOpenIDRequestValidator


def _std_hash_by_alg(alg) -> Callable[[bytes], Any]:
    if alg.endswith("S256"):
        return hashlib.sha256
    if alg.endswith("S384"):
        return hashlib.sha384
    if alg.endswith("S512"):
        return hashlib.sha512
    raise ValueError("{} is not supported".format(alg))


def _left_half_hash(data: str, alg: str) -> str:
    hashfun = _std_hash_by_alg(alg)
    hashobj = hashfun(data.encode("ascii"))
    digest = hashobj.digest()[: hashobj.digest_size // 2]
    return b64encode(digest)


class OpenIDConnectGrant(
    AbstractBaseResponseTypeHandler, AbstractBaseGrantTypeHandler
):
    """
    http://openid.net/specs/openid-connect-core-1_0.html
    """

    __slots__ = ["code_factory", "token_factory", "jwt_encoder"]
    jwt_encoder_cls = JWT

    def __init__(
        self,
        code_factory: Callable[
            [Any, str, str, OAuth2RequestMixin], Dict[str, Any]
        ] = authorization_code_factory,
        token_factory: Callable[
            [Any, str, bool, OAuth2RequestMixin], Dict[str, Any]
        ] = bearer_token_factory,
    ) -> None:
        self.code_factory = code_factory
        self.token_factory = token_factory
        self.jwt_encoder = self.jwt_encoder_cls()

    def supported_response_types(self) -> List[str]:
        return [
            # Authorization Code Flow
            "code",
            # Implicit Flow
            "id_token",
            "id_token token",
            # Hybrid Flow
            "code id_token",
            "code token",
            "code id_token token",
            # Multiple Response Type Encoding Practices
            "none",
            # Oops
            "token",
        ]

    def get_default_response_mode(
        self, request: OpenIDConnectRequestMixin
    ) -> str:
        if request.o2_response_type in ("code", "none"):
            return "query"
        return "fragment"

    def is_response_mode_allowed(
        self, response_mode: str, request: OpenIDConnectRequestMixin
    ) -> bool:
        """
        http://openid.net/specs/oauth-v2-multiple-response-types-1_0.html#Combinations
        """
        if request.o2_response_type in ("code", "none"):
            return response_mode in ("query", "fragment")
        return response_mode == "fragment"

    def validate_authorization_request(
        self,
        request: OpenIDConnectRequestMixin,
        repository: AbstractOpenIDRepository,
        validator: AbstractOpenIDRequestValidator,
    ) -> None:
        if (
            not request.o2_validate(requireds={"scope"})
            or "openid" not in request.o2_scope.split()
        ):
            # the request is not for Open ID Connect.
            AbstractBaseResponseTypeHandler.validate_authorization_request(
                self, request, repository, validator
            )
            return

        if not request.o2_validate(requireds={"client_id", "redirect_uri"}):
            raise FatalOAuth2Error()
        if not repository.is_client_available(request.o2_client_id, request):
            raise FatalOAuth2Error()
        # NOTE(yosida95): 'response_type' parameter was eventually validated by
        # the AuthorizationEndpoint to dispatch this method, the value would be
        # one of the self.supported_response_types(), so the value can be used
        # safely here.
        resp_type = request.o2_response_type.split()
        if not request.o2_validate_redirect_uri(
            repository, validator, exactly="code" not in resp_type
        ):
            raise FatalOAuth2Error()

        if not request.o2_validate(
            optionals={
                "state",
                "nonce",
                "display",
                "prompt",
                "max_age",
                "ui_locales",
                "id_token_hint",
                "login_hint",
                "acr_values",
            }
        ):
            raise InvalidRequestError(request)
        if "code" not in resp_type and not request.o2_nonce:
            # In implicit flow, use of nonce is REQUIRED.
            raise InvalidRequestError(request)

        if not validator.validate_scope(
            request.o2_scope, request.o2_client_id, request
        ):
            raise InvalidScopeError(request)

    def handle_authorization_request(
        self,
        request: OpenIDConnectRequestMixin,
        repository: AbstractOpenIDRepository,
        subject: Any,
    ) -> None:
        resp = dict()
        if request.o2_state:
            resp["state"] = request.o2_state

        try:
            resp_type = request.o2_response_type.split()

            code = None
            if "code" in resp_type:
                code = self.code_factory(
                    subject, request.o2_scope, request.o2_client_id, request
                )
                repository.save_authorization_code(
                    code,
                    request.o2_redirect_uri,
                    request.o2_client_id,
                    request,
                )
                resp.update(code.copy())

            token = None
            if "token" in resp_type:
                with_refresh_token = False
                token = self.token_factory(
                    subject, request.o2_scope, with_refresh_token, request
                )
                repository.save_access_token(token, request)
                resp.update(token.copy())

            if "id_token" in resp_type:
                key, alg = repository.get_key_alg_by_client_id(
                    request.o2_client_id, request
                )

                jwt_dict = repository.get_id_token_by_subject(
                    subject, code, token, request
                )
                if request.o2_nonce:
                    jwt_dict["nonce"] = request.o2_nonce
                if code:
                    jwt_dict["c_hash"] = _left_half_hash(code["code"], alg)
                if token:
                    jwt_dict["at_hash"] = _left_half_hash(
                        token["access_token"], alg
                    )
                if code is None and token is None:
                    userinfo = repository.get_userinfo_by_subject(
                        subject, request.o2_scope, request
                    )
                    jwt_dict = jwt_dict.copy()
                    if jwt_dict["sub"] != userinfo["sub"]:
                        raise ServerError(request)
                    jwt_dict.update(userinfo)
                resp["id_token"] = self.jwt_encoder.encode(
                    jwt_dict, key=key, alg=alg
                )
        except BaseException as why:
            raise ServerError(request) from why

        redirect_uri = extend_uri_with_params(
            request.o2_redirect_uri, resp, request.o2_response_mode
        )
        request.response.state = 302
        request.response.location = redirect_uri
        return None

    def supported_grant_types(self) -> List[str]:
        return ["authorization_code"]

    def validate_token_request(
        self,
        request: OpenIDConnectRequestMixin,
        repository: AbstractOpenIDRepository,
        validator: AbstractOpenIDRequestValidator,
    ) -> None:
        super().validate_token_request(request, repository, validator)

        if not request.o2_validate(
            requireds={"code"}, optionals={"redirect_uri"}
        ):
            raise InvalidRequestError(request)

        if not request.o2_validate_redirect_uri(repository, validator):
            raise InvalidGrantError(request)

        if not validator.validate_authorization_code(
            request.o2_code,
            request.o2_redirect_uri,
            request.o2_client_id,
            request,
        ):
            raise InvalidGrantError(request)

    def handle_token_request(
        self,
        request: OpenIDConnectRequestMixin,
        repository: AbstractOpenIDRepository,
    ) -> dict:
        resp = {}
        try:
            subject = repository.get_subject_by_authorization_code(
                request.o2_code, request
            )
            scope = repository.get_scope_by_authorization_code(
                request.o2_code, request
            )
            token = self.token_factory(subject, scope, True, request)
            repository.save_access_token(token, request)
            resp.update(token.copy())

            key, alg = repository.get_key_alg_by_client_id(
                request.o2_client_id, request
            )
            jwt_dict = repository.get_id_token_by_subject(
                subject, {"code": request.o2_code}, token, request
            )
            jwt_dict["at_hash"] = _left_half_hash(token["access_token"], alg)
            resp["id_token"] = self.jwt_encoder.encode(
                jwt_dict, key=key, alg=alg
            )
        except BaseException as why:
            raise ServerError(request) from why

        return resp
