from authlib.oauth2.rfc6749.grants import (
    AuthorizationCodeGrant as _AuthorizationCodeGrant,
)
from authlib.oidc.core import OpenIDCode as OpenIDCode_
from authlib.oidc.core import UserInfo
from aws_lambda_powertools.event_handler.exceptions import NotFoundError
from layercake.dateutils import now, ttl
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair

from apigateway_oauth2 import (
    AuthorizationServer,
    OAuth2Client,
    OAuth2Token,
)
from boto3clients import dynamodb_client
from config import DYNAMODB_SORT_KEY, OAUTH2_TABLE

oauth2_layer = DynamoDBPersistenceLayer(OAUTH2_TABLE, dynamodb_client)


DUMMY_JWT_CONFIG = {
    'key': 'secret-key',
    'alg': 'HS256',
    'iss': 'https://authlib.org',
    'exp': 3600,
}


def create_save_token_func(persistence_layer: DynamoDBPersistenceLayer):
    def save_token(token, request) -> OAuth2Token:
        return OAuth2Token()

    return save_token


class ClientNotFoundError(NotFoundError):
    def __init__(self, *_):
        super().__init__('Client not found')


def create_query_client_func(persistence_layer: DynamoDBPersistenceLayer):
    def query_client(client_id) -> OAuth2Client:
        client = persistence_layer.collection.get_item(
            KeyPair('OAUTH2_CLIENT', f'CLIENT_ID#{client_id}'),
            exc_cls=ClientNotFoundError,
        )

        _, client_id = client.get(DYNAMODB_SORT_KEY, '').split('#')

        return OAuth2Client(
            client_id=client_id,
            redirect_uris=client['redirect_uris'],
            response_types=client['response_types'],
        )

    return query_client


def save_authorization_code(code, request):
    data: dict = request.payload.data  # type: ignore
    user: dict = request.user  # type: ignore
    nonce: str | None = data.get('nonce')
    now_ = now()
    ttl_ = ttl(start_dt=now_, minutes=15)

    with oauth2_layer.transact_writer() as transact:
        transact.put(
            item={
                'id': f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
                'sk': f'CODE#{code}',
                'redirect_uri': request.payload.redirect_uri,  # type: ignore
                'scope': request.payload.scope,  # type: ignore
                'user_id': user['id'],
                'nonce': nonce,
                'created_at': now_,
                'ttl': ttl_,
            },
        )

        if nonce:
            transact.put(
                item={
                    'id': f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
                    'sk': f'NONCE#{nonce}',
                    'code': code,
                    'created_at': now_,
                    'ttl': ttl_,
                },
            )


def exists_nonce(nonce, request):
    nonce_ = oauth2_layer.get_item(
        KeyPair(
            f'OAUTH2_CODE#CLIENT_ID#{request.payload.client_id}',
            f'NONCE#{nonce}',
        )
    )
    return bool(nonce_)


class OpenIDCode(OpenIDCode_):
    def exists_nonce(self, nonce, request):
        return exists_nonce(nonce, request)

    def get_jwt_config(self, grant):
        return DUMMY_JWT_CONFIG

    def generate_user_info(self, user, scope):
        return UserInfo(
            sub=user.id,
            name=user.name,
            email=user.email,
        ).filter(scope)


class AuthorizationCodeGrant(_AuthorizationCodeGrant):
    TOKEN_ENDPOINT_AUTH_METHODS = [
        'client_secret_basic',
        'client_secret_post',
    ]

    def save_authorization_code(self, code: str, request):
        return save_authorization_code(code, request)


authorization = AuthorizationServer(
    query_client=create_query_client_func(oauth2_layer),
    save_token=create_save_token_func(oauth2_layer),
)
authorization.register_grant(AuthorizationCodeGrant, [OpenIDCode(require_nonce=False)])