from abc import ABC
from dataclasses import dataclass
from datetime import timedelta
from enum import Enum
from typing import NotRequired, TypedDict

from layercake.dateutils import now, ttl
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
from layercake.strutils import md5_hash

from schemas import Enrollment

Org = TypedDict(
    'Org',
    {
        'org_id': str,
        'name': str,
    },
)

DeduplicationWindow = TypedDict(
    'DeduplicationWindow',
    {
        'offset_days': int,
    },
)

Subscription = TypedDict(
    'Subscription',
    {
        'org_id': str,
        'billing_day': int,
        'billing_period': NotRequired[str],
    },
)


class Kind(str, Enum):
    ORDER = 'ORDER'
    ENROLLMENT = 'ENROLLMENT'


@dataclass(frozen=True)
class LinkedEntity(ABC):
    id: str
    kind: Kind
    table_name: str | None = None


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


def enroll(
    enrollment: Enrollment,
    *,
    org: Org | None = None,
    subscription: Subscription | None = None,
    linked_entities: frozenset[LinkedEntity] = frozenset(),
    deduplication_window: DeduplicationWindow | None = None,
    persistence_layer: DynamoDBPersistenceLayer,
) -> bool:
    now_ = now()
    user = enrollment.user
    course = enrollment.course
    lock_hash = md5_hash('%s%s' % (user.id, course.id))
    access_expires_at = now_ + timedelta(days=course.access_period)

    with persistence_layer.transact_writer() as transact:
        transact.put(
            item={
                'sk': '0',
                'created_at': now_,
                'access_expires_at': access_expires_at,
                **enrollment.model_dump(),
            }
            | ({'subscription_covered': True} if subscription else {})
            | ({'org_id': org['org_id']} if org else {}),
        )

        # Relationships between this enrollment and its related entities
        for entity in linked_entities:
            # Parent knows the child
            transact.put(
                item={
                    'id': entity.id,
                    'sk': f'LINKED_ENTITIES#{entity.kind.value}#CHILD',
                    'created_at': now_,
                    'enrollment_id': enrollment.id,
                },
                cond_expr='attribute_not_exists(sk)',
                table_name=entity.table_name,
            )

            keyprefix = entity.kind.value.lower()
            # Child knows the parent
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': f'LINKED_ENTITIES#{entity.kind.value}#PARENT',
                    'created_at': now_,
                    f'{keyprefix}_id': entity.id,
                },
                cond_expr='attribute_not_exists(sk)',
            )

        if org:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'ORG',
                    'created_at': now_,
                }
                | org
            )

        if subscription:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#SUBSCRIPTION_COVERED',
                    'created_at': now_,
                }
                | subscription,
            )

        # Prevents the user from enrolling in the same course again until
        # the deduplication window expires or is removed.
        if deduplication_window:
            offset_days = int(deduplication_window['offset_days'])
            ttl_ = 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_,
                },
                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_,
                },
            )
            # 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