from dataclasses import asdict, dataclass
from typing import Self, TypedDict
from uuid import uuid4

from aws_lambda_powertools.event_handler.exceptions import (
    BadRequestError,
    NotFoundError,
)
from layercake.dateutils import now, ttl
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
from layercake.strutils import md5_hash

from config import ORDER_TABLE
from models import Course, Enrollment

Tenant = TypedDict('Tenant', {'id': str, 'name': str})
Author = TypedDict('Author', {'id': str, 'name': str})
DeduplicationWindow = TypedDict('DeduplicationWindow', {'offset_days': int})


class LinkedEntity(str):
    def __new__(cls, id: str, type: str) -> Self:
        return super().__new__(cls, '#'.join([type.upper(), id]))

    def __init__(self, id: str, type: str) -> None:
        # __init__ is used to store the parameters for later reference.
        # For immutable types like str, __init__ cannot change the instance's value.
        self.id = id
        self.type = type


@dataclass(frozen=True)
class Slot:
    id: str
    sk: str

    @property
    def order_id(self) -> LinkedEntity:
        idx, _ = self.sk.split('#')
        return LinkedEntity(idx, 'order')


class SlotDoesNotExistError(NotFoundError):
    def __init__(self, *args):
        super().__init__('Slot not found')


class DeduplicationConflictError(BadRequestError):
    def __init__(self, *args):
        super().__init__('Enrollment already exists')


class EnrollmentConflictError(BadRequestError):
    def __init__(self, *_):
        super().__init__('Enrollment status conflict')


def enroll(
    enrollment: Enrollment,
    *,
    tenant: Tenant,
    slot: Slot | None = None,
    author: Author | None = None,
    linked_entities: frozenset[LinkedEntity] = frozenset(),
    deduplication_window: DeduplicationWindow | None = None,
    persistence_layer: DynamoDBPersistenceLayer,
) -> bool:
    """Enrolls a user into a course and schedules lifecycle events."""
    now_ = now()
    user = enrollment.user
    course = enrollment.course
    org_id = tenant['id']
    lock_hash = md5_hash('%s%s' % (user.id, course.id))

    with persistence_layer.transact_writer() as transact:
        if slot:
            linked_entities = frozenset({slot.order_id}) | linked_entities

        transact.put(
            item={
                'sk': '0',
                'created_at': now_,
                'org_id': org_id,
                **enrollment.model_dump(),
            },
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'ORG',
                'org_id': org_id,
                'name': tenant['name'],
                'created_at': now_,
            },
        )

        for entity in linked_entities:
            type = entity.type.lower()
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': f'LINKED_ENTITIES#{type}',
                    'created_at': now_,
                    f'{type}_id': entity.id,
                }
            )

        if slot:
            transact.put(
                item={
                    'id': enrollment.id,
                    # Post-migration: uncomment the following line
                    # 'sk': 'metadata#parent_slot',
                    'sk': 'parent_vacancy',
                    'vacancy': asdict(slot),
                    'created_at': now_,
                }
            )

            transact.delete(
                key=KeyPair(slot.id, slot.sk),
                cond_expr='attribute_exists(sk)',
                exc_cls=SlotDoesNotExistError,
            )
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'CANCEL_POLICY',
                    'created_at': now_,
                }
            )

        if author:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'author',
                    'user_id': author['id'],
                    'name': author['name'],
                    'created_at': now_,
                },
            )

        # Prevents the user from enrolling in the same course again until
        # the deduplication window expires or is removed.
        if deduplication_window:
            offset_days = deduplication_window['offset_days']
            ttl_expiration = ttl(start_dt=now_, days=course.access_period - offset_days)
            transact.put(
                item={
                    'id': 'LOCK',
                    'sk': lock_hash,
                    'enrollment_id': enrollment.id,
                    'created_at': now_,
                    'ttl': ttl_expiration,
                },
                cond_expr='attribute_not_exists(sk)',
                exc_cls=DeduplicationConflictError,
            )
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'LOCK',
                    'hash': lock_hash,
                    'created_at': now_,
                    'ttl': ttl_expiration,
                },
            )
            # Deduplication window can be recalculated if needed
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#DEDUPLICATION_WINDOW',
                    'offset_days': offset_days,
                    'created_at': now_,
                },
            )
        else:
            transact.condition(
                key=KeyPair('LOCK', lock_hash),
                cond_expr='attribute_not_exists(sk)',
                exc_cls=DeduplicationConflictError,
            )

    return True


def set_status_as_canceled(
    id: str,
    *,
    lock_hash: str | None = None,
    created_by: Author,
    course: Course | None = None,
    vacancy_key: KeyPair | None = None,
    persistence_layer: DynamoDBPersistenceLayer,
):
    """Cancel the enrollment if there's a `cancel_policy`
    and put its vacancy back if `vacancy_key` is provided."""
    now_ = now()

    with persistence_layer.transact_writer() as transact:
        transact.update(
            key=KeyPair(id, '0'),
            update_expr='SET #status = :canceled, \
                             access_expired = :true, \
                             canceled_at = :now, \
                             updated_at = :now',
            cond_expr='#status = :pending',
            expr_attr_names={
                '#status': 'status',
            },
            expr_attr_values={
                ':canceled': 'CANCELED',
                ':pending': 'PENDING',
                ':true': True,
                ':now': now_,
            },
        )
        transact.put(
            item={
                'id': id,
                'sk': 'CANCELED',
                'canceled_by': created_by,
                'created_at': now_,
            },
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='CANCEL_POLICY',
            ),
            cond_expr='attribute_exists(sk)',
            exc_cls=EnrollmentConflictError,
        )
        # Remove reminders and policies that no longer apply
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='SCHEDULE#REMINDER_NO_ACCESS_AFTER_3_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='SCHEDULE#REMINDER_NO_ACTIVITY_AFTER_7_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='SCHEDULE#REMINDER_ACCESS_PERIOD_BEFORE_30_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='SCHEDULE#REMINDER_CERT_EXPIRATION_BEFORE_30_DAYS',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='SCHEDULE#SET_ACCESS_EXPIRED',
            )
        )
        transact.delete(
            key=KeyPair(
                pk=id,
                sk='parent_vacancy',
            )
        )

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

        if vacancy_key and course:
            vacancy_pk, vacancy_sk = vacancy_key.values()
            org_id = vacancy_pk.removeprefix('vacancies#')
            order_id, enrollment_id = vacancy_sk.split('#')

            transact.condition(
                key=KeyPair(order_id, '0'),
                cond_expr='attribute_exists(id)',
                table_name=ORDER_TABLE,
            )
            # Put the vacancy back and assign a new ID
            transact.put(
                item={
                    # Post-migration: uncomment the following line
                    # 'id': f'slots#org#{org_id}',
                    'id': f'vacancies#{org_id}',
                    'sk': f'{order_id}#{uuid4()}',
                    'course': course,
                    'created_at': now_,
                },
                cond_expr='attribute_not_exists(sk)',
            )
            # Set the status of `generated_items` to `ROLLBACK` to know
            # which slot is available for reuse
            transact.update(
                # Post-migration: uncomment the following line
                # key=KeyPair(order_id, f'slots#enrollment#{enrollment_id}'),
                key=KeyPair(order_id, f'generated_items#{enrollment_id}'),
                update_expr='SET #status = :status, updated_at = :updated_at',
                expr_attr_names={
                    '#status': 'status',
                },
                expr_attr_values={
                    ':status': 'ROLLBACK',
                    ':updated_at': now_,
                },
                cond_expr='attribute_exists(sk)',
                table_name=ORDER_TABLE,
            )

    return True
