from typing import Annotated

from aws_lambda_powertools import Logger
from aws_lambda_powertools.event_handler.api_gateway import Router
from aws_lambda_powertools.event_handler.exceptions import (
    BadRequestError,
)
from aws_lambda_powertools.event_handler.openapi.params import Body
from layercake.dateutils import now
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair

from boto3clients import dynamodb_client
from config import ENROLLMENT_TABLE

logger = Logger(__name__)
router = Router()
dyn = DynamoDBPersistenceLayer(ENROLLMENT_TABLE, dynamodb_client)


class CancelPolicyConflictError(BadRequestError):
    def __init__(self, *_):
        super().__init__('Cancellation policy not found')


class SlotConflictError(BadRequestError):
    def __init__(self, *_):
        super().__init__('Slot not found')


@router.post('/<enrollment_id>/cancel')
def cancel(
    enrollment_id: str,
    lock_hash: Annotated[str | None, Body(embed=True)] = None,
):
    now_ = now()

    with dyn.transact_writer() as transact:
        transact.update(
            key=KeyPair(enrollment_id, '0'),
            cond_expr='#status = pending',
            update_expr='SET #status = :canceled, \
                             canceled_at = :now, \
                             updated_at = :now',
            expr_attr_names={
                ':status': 'status',
            },
            expr_attr_values={
                ':pending': 'PENDING',
                ':canceled': 'CANCELED',
                ':true': True,
                ':now': now_,
            },
        )
        transact.put(
            item={
                'id': enrollment_id,
                'sk': 'CANCELED_BY',
                'canceled_by': {},
                'created_at': now_,
            }
        )
        transact.delete(
            key=KeyPair(
                pk=enrollment_id,
                sk='CANCEL_POLICY',
            ),
            cond_expr='attribute_exists(sk)',
            exc_cls=CancelPolicyConflictError,
        )
        # Remove reminders and policies that no longer apply
        transact.delete(
            key=KeyPair(
                pk=enrollment_id,
                sk='SCHEDULE#REMINDER_NO_ACCESS_AFTER_3_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=enrollment_id,
                sk='SCHEDULE#REMINDER_ACCESS_PERIOD_BEFORE_30_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=enrollment_id,
                sk='METADATA#PARENT_SLOT',
            ),
            cond_expr='attribute_exists(sk)',
            exc_cls=SlotConflictError,
        )

        if lock_hash:
            transact.delete(key=KeyPair(enrollment_id, 'LOCK'))
            transact.delete(key=KeyPair('LOCK', lock_hash))
