from datetime import timedelta
from typing import NotRequired, Self, 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 LinkedEntity(str):
    def __new__(cls, id: str, type: str) -> Self:
        return super().__new__(cls, '#'.join([type.lower(), 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


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:
    """Enrolls a user into a course and schedules lifecycle events."""
    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 {})
            | ({'tenant_id': org['org_id']} if org else {}),
            # Post-migration: uncomment the following line
            # | ({'org_id': org['org_id']} if org else {}),
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'METADATA#COURSE',
                'created_at': now_,
                **course.model_dump(include={'cert', 'access_period'}),
            }
        )

        # Relationships between this enrollment and its related entities
        for parent_entity in linked_entities:
            perent_id = parent_entity.id
            entity_sk = f'LINKED_ENTITIES#{parent_entity.type}'
            keyprefix = parent_entity.type.lower()

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

            # Parent knows the child
            transact.put(
                item={
                    'id': perent_id,
                    'sk': f'{entity_sk}#CHILD',
                    'created_at': now_,
                    f'{keyprefix}_id': enrollment.id,
                },
                cond_expr='attribute_not_exists(sk)',
            )

        if org:
            transact.put(
                item={
                    'id': enrollment.id,
                    #  Post-migration: uncomment the following line
                    # 'sk': 'ORG',
                    'sk': 'tenant',
                    '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