# Copyright (C) 2023 Cochise Ruhulessin
#
# All rights reserved. No warranty, explicit or implicit, provided. In
# no event shall the author(s) be liable for any claim or damages.
#
# 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.
from typing import Awaitable
from typing import Callable
from typing import TypeVar
from typing import Union

import fastapi
import pydantic
from headless.ext.oauth2 import Client
from headless.ext.oauth2 import TokenResponse

from cbra.types import IDependant
from .authorizationstate import AuthorizationState
from .authorizationstateidentifier import AuthorizationStateIdentifier
from .authorizecoderesponse import AuthorizeCodeResponse
from .authorizeerrorresponse import AuthorizeErrorResponse
from .authorizetokenresponse import AuthorizeTokenResponse
from .fatalauthorizationexception import FatalAuthorizationException
from .ifrontendstorage import IFrontendStorage


T = TypeVar('T', bound='AuthorizeResponse')


class AuthorizeResponse(pydantic.BaseModel, IDependant):
    __root__: Union[
        AuthorizeCodeResponse,
        AuthorizeErrorResponse,
        AuthorizeTokenResponse
    ]

    @classmethod
    def __inject__(cls: type[T]) -> Callable[..., Awaitable[T] | T]:
        return cls.parse_query

    @classmethod
    async def parse_query(
        cls: type[T],
        code: str | None = fastapi.Query(
            default=None,
            title='Authorization code',
            description=(
                'The authorization code generated by the authorization server. '
                'Required if using the Authorization Code Flow and the response '
                'mode is `query`, otherwise this parameter is ignored.'
            )
        ),
        state: AuthorizationStateIdentifier | None = fastapi.Query(
            default=None,
            title='State',
            description=(
                '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. If supplied, this '
                'parameter **must** equal the `state` parameter used when '
                'creating the authorization request. The `state` parameter '
                'is ignored when using JARM because it is included in the '
                'JSON Web Token supplied using the `jwt` parameter per '
                'chosen response mode.'
            )
        ),
        error: str | None = fastapi.Query(
            default=None,
            title="Error code",
            description=(
                "The error code returned by the authorization server if "
                "the user cancelled the request, refused consent, or "
                "failed to authenticate."
            )
        ),
        error_description: str | None = fastapi.Query(
            default=None,
            title="Error description",
            description=(
                "Human-readable ASCII [USASCII] text providing "
                "additional information, used to assist the client "
                "developer in understanding the error that occurred."
            )
        ),
        error_uri: str | None = fastapi.Query(
            default=None,
            title="Error URI",
            description=(
                "A URI identifying a human-readable web page with information "
                "about the error, used to provide the client developer with "
                "additional information about the error."
            )
        ),
        iss: str | None = fastapi.Query(
            default=None,
            title='Issuer',
            description=(
                'Identifies the authorization server that redirected to '
                'this endpoint, as defined in RFC 9207. If the client supports '
                'OAuth 2.0 Authorization Server Issuer Identification, then '
                'this parameter is **required**, if the response mode is not '
                'JWT Secured Authorization Response Mode (JARM) per RFC 9101.\n\n'
                'The `iss` parameter value is the issuer identifier of '
                'the authorization server that created the authorization '
                'response, as defined in RFC 8414.  Its value **must** '
                'be a URL that uses the `https` scheme without any '
                'query or fragment components.'
            )
        ),
        response: str | None = fastapi.Query(
            default=None,
            title="Response",
            alias='response',
            description=(
                "If the response mode is `query.jwt`, then this parameter contains the JSON "
                "Web Token (JWT) containing the response from the authorization server. The "
                "JWT is signed first by the authorization server, then optionally encrypted "
                "using a pre-registered public key that is associated to the client."
            )
        )
    ) -> T:
        if iss:
            raise NotImplementedError("RFC 9207 is not implemented")
        self = cls.parse_obj({
            'code': code,
            'error': error,
            'error_description': error_description,
            'error_uri': error_uri,
            'iss': iss,
            'response': response,
            'state': state
        })
        if isinstance(self.__root__, AuthorizeTokenResponse):
            raise NotImplementedError(
                "JWT Secured Authorization Response Mode is not implemented."
            )
        return self
    
    async def decrypt(self, client: Client) -> None:
        raise NotImplementedError

    async def obtain(self, client: Client, params: AuthorizationState) -> TokenResponse:
        assert isinstance(self.__root__, AuthorizeCodeResponse)
        return await self.__root__.obtain(client, params)

    async def verify(
        self,
        client: Client,
        storage: IFrontendStorage
    ) -> AuthorizationState:
        """Verifies that the authorization response is issued by the authorization
        server and that the response parameters are correct."""
        if isinstance(self.__root__, AuthorizeErrorResponse):
            self.__root__.raise_exception()
            assert False
        if isinstance(self.__root__, AuthorizeTokenResponse):
            await self.decrypt(client)
            assert isinstance(self.__root__, AuthorizeCodeResponse)
        params = await storage.get(self.get_state())
        if params is None:
            raise FatalAuthorizationException(
                "No authorization request satisfied the query parameters."
            )
        return params

    def get_state(self) -> AuthorizationStateIdentifier:
        assert isinstance(self.__root__, AuthorizeCodeResponse)
        return self.__root__.state