from datetime import date, datetime, time, timedelta
from typing import Annotated, NotRequired, TypedDict
from uuid import uuid4

import pytz
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,
    NotFoundError,
)
from aws_lambda_powertools.event_handler.openapi.params import Body
from aws_lambda_powertools.shared.functions import extract_event_from_common_models
from layercake.batch import BatchProcessor
from layercake.dateutils import now, ttl
from layercake.dynamodb import DynamoDBPersistenceLayer, KeyPair
from layercake.extra_types import CpfStr, NameStr
from layercake.strutils import md5_hash
from pydantic import UUID4, BaseModel, EmailStr, Field, FutureDate

from boto3clients import dynamodb_client
from config import (
    DEDUP_WINDOW_OFFSET_DAYS,
    ENROLLMENT_TABLE,
    ORDER_TABLE,
    TZ,
    USER_TABLE,
)
from exceptions import (
    ConflictError,
    OrderNotFoundError,
    SubscriptionConflictError,
    SubscriptionFrozenError,
    SubscriptionRequiredError,
)
from middlewares.authentication_middleware import User as Authenticated

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


class DeduplicationConflictError(ConflictError): ...


class SeatNotFoundError(NotFoundError): ...


class User(BaseModel):
    id: str | UUID4
    name: NameStr
    email: EmailStr
    cpf: CpfStr


class Course(BaseModel):
    id: UUID4
    name: str
    access_period: int


class DeduplicationWindow(BaseModel):
    offset_days: int


class Subscription(BaseModel):
    billing_day: int


class Seat(BaseModel):
    order_id: UUID4


class Enrollment(BaseModel):
    id: UUID4 = Field(default_factory=uuid4)
    user: User
    course: Course
    scheduled_for: FutureDate | None = None
    deduplication_window: DeduplicationWindow | None = None
    seat: Seat | None = None


class Org(BaseModel):
    id: str | UUID4
    name: str


@router.post('/')
def enroll(
    org_id: Annotated[str | UUID4, Body(embed=True)],
    enrollments: Annotated[tuple[Enrollment, ...], Body(embed=True)],
    subscription: Annotated[Subscription | None, Body(embed=True)] = None,
):
    now_ = now()
    created_by: Authenticated = router.context['user']
    org = dyn.collection.get_item(
        KeyPair(
            pk=str(org_id),
            sk='0',
            table_name=USER_TABLE,
        )
    )

    ctx = {
        'org': Org.model_validate(org),
        'created_by': created_by,
        'subscription': subscription,
    }

    immediate = [e for e in enrollments if not e.scheduled_for]
    later = [e for e in enrollments if e.scheduled_for]

    with processor(immediate, enroll_now, ctx) as batch:
        now_out = batch.process()

    with processor(later, _enroll_later, ctx) as batch:
        later_out = batch.process()

    def fmt(r):
        return {
            'status': r.status.value,
            'input_record': extract_event_from_common_models(r.input_record),
            'output': extract_event_from_common_models(r.output),
            'cause': r.cause,
        }

    item = {
        'id': f'SUBMISSION#ORG#{org_id}',
        'sk': now_,
        'enrolled': list(map(fmt, now_out)) if now_out else None,
        'scheduled': list(map(fmt, later_out)) if later_out else None,
        'ttl': ttl(start_dt=now_, days=30 * 3),
        'created_by': {
            'id': created_by.id,
            'name': created_by.name,
        },
    }

    try:
        dyn.put_item(item=item)
    except Exception as exc:
        logger.exception(exc)
    else:
        return item


Context = TypedDict(
    'Context',
    {
        'org': Org,
        'created_by': Authenticated,
        'subscription': NotRequired[Subscription],
    },
)


def enroll_now(enrollment: Enrollment, context: Context):
    now_ = now()
    user = enrollment.user
    course = enrollment.course
    seat = enrollment.seat
    org = context['org']
    subscription = context.get('subscription')
    created_by = context['created_by']
    lock_hash = md5_hash(f'{user.id}{course.id}')
    access_expires_at = now_ + timedelta(days=course.access_period)
    deduplication_window = enrollment.deduplication_window
    offset_days = (
        int(deduplication_window.offset_days)
        if deduplication_window
        else DEDUP_WINDOW_OFFSET_DAYS
    )
    dedup_lock_ttl = ttl(
        start_dt=now_,
        days=course.access_period - offset_days,
    )

    if not (bool(subscription) ^ bool(seat)):
        raise BadRequestError('Malformed body')

    with dyn.transact_writer() as transact:
        transact.put(
            item={
                'id': enrollment.id,
                'sk': '0',
                'score': None,
                'progress': 0,
                'status': 'PENDING',
                'user': user.model_dump(),
                'course': course.model_dump(),
                'access_expires_at': access_expires_at,
                'org_id': org.id,
                'created_at': now_,
            }
            | ({'subscription_covered': True} if subscription else {})
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'ORG',
                'name': org.name,
                'org_id': org.id,
                'created_at': now_,
            }
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'CANCEL_POLICY',
                'created_at': now_,
            }
            | ({'seat': seat.model_dump()} if seat else {})
        )

        if seat:
            transact.condition(
                key=KeyPair(str(seat.order_id), '0'),
                cond_expr='attribute_exists(sk)',
                exc_cls=OrderNotFoundError,
                table_name=ORDER_TABLE,
            )
            transact.update(
                key=KeyPair(
                    pk=str(seat.order_id),
                    sk=f'ENROLLMENT#{enrollment.id}',
                ),
                update_expr='SET course = :course, \
                                 #user = :user, \
                                 #status = :executed, \
                                 executed_at = :now, \
                                 created_at = if_not_exists(created_at, :now)',
                expr_attr_names={
                    '#user': 'user',
                    '#status': 'status',
                },
                expr_attr_values={
                    ':course': course.model_dump(),
                    ':user': user.model_dump(),
                    ':executed': 'EXECUTED',
                    ':now': now_,
                },
                table_name=ORDER_TABLE,
            )
            transact.delete(
                key=KeyPair(
                    f'SEAT#ORG#{org.id}',
                    f'ORDER#{seat.order_id}#ENROLLMENT#{enrollment.id}',
                ),
                cond_expr='attribute_exists(sk)',
                exc_cls=SeatNotFoundError,
            )
            # Enrollment should know where it comes from
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': f'LINKED_ENTITY#PARENT#ORDER#{seat.order_id}',
                    'created_at': now_,
                },
                cond_expr='attribute_not_exists(sk)',
            )

        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'CREATED_BY',
                'name': created_by.name,
                'user_id': created_by.id,
                'created_at': now_,
            }
        )
        transact.put(
            item={
                'id': enrollment.id,
                'sk': 'LOCK',
                'hash': lock_hash,
                'created_at': now_,
                'ttl': dedup_lock_ttl,
            },
        )
        transact.put(
            item={
                'id': 'LOCK',
                'sk': lock_hash,
                'enrollment_id': enrollment.id,
                'created_at': now_,
                'ttl': dedup_lock_ttl,
            },
            cond_expr='attribute_not_exists(sk)',
            exc_cls=DeduplicationConflictError,
        )

        if subscription:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#SUBSCRIPTION_COVERED',
                    'org_id': org.id,
                    'billing_day': subscription.billing_day,
                    'created_at': now_,
                }
            )
            transact.condition(
                key=KeyPair('SUBSCRIPTION', f'ORG#{org.id}'),
                cond_expr='attribute_exists(sk)',
                exc_cls=SubscriptionRequiredError,
                table_name=USER_TABLE,
            )
            transact.condition(
                key=KeyPair(str(org.id), 'METADATA#SUBSCRIPTION'),
                cond_expr='billing_day = :billing_day',
                expr_attr_values={
                    ':billing_day': subscription.billing_day,
                },
                exc_cls=SubscriptionConflictError,
                table_name=USER_TABLE,
            )
            transact.condition(
                key=KeyPair('SUBSCRIPTION#FROZEN', f'ORG#{org.id}'),
                cond_expr='attribute_not_exists(sk)',
                exc_cls=SubscriptionFrozenError,
                table_name=USER_TABLE,
            )

        # The deduplication window can be recalculated based on user settings.
        if deduplication_window:
            transact.put(
                item={
                    'id': enrollment.id,
                    'sk': 'METADATA#DEDUPLICATION_WINDOW',
                    'offset_days': offset_days,
                    'created_at': now_,
                },
            )

    return enrollment


def _enroll_later(enrollment: Enrollment, context: Context):
    now_ = now()
    user = enrollment.user
    course = enrollment.course
    seat = enrollment.seat
    scheduled_for = _date_to_midnight(enrollment.scheduled_for)  # type: ignore
    dedup_window = enrollment.deduplication_window
    org = context['org']
    subscription = context.get('subscription')
    created_by = context['created_by']
    lock_hash = md5_hash(f'{user.id}{course.id}')

    if not (bool(subscription) ^ bool(seat)):
        raise BadRequestError('Malformed body')

    with dyn.transact_writer() as transact:
        pk = f'SCHEDULED#ORG#{org.id}'
        sk = f'{scheduled_for.isoformat()}#{lock_hash}'

        transact.put(
            item={
                'id': pk,
                'sk': sk,
                'user': user.model_dump(),
                'course': course.model_dump(),
                'org_name': org.name,
                'enrollment_id': enrollment.id,
                'created_by': {
                    'id': created_by.id,
                    'name': created_by.name,
                },
                'ttl': ttl(start_dt=scheduled_for),
                'scheduled_at': now_,
            }
            | ({'seat': seat.model_dump()} if seat else {})
            | (
                {'dedup_window_offset_days': dedup_window.offset_days}
                if dedup_window
                else {}
            )
            | (
                {
                    'subscription_billing_day': subscription.billing_day,
                }
                if subscription
                else {}
            ),
        )

        if seat:
            transact.condition(
                key=KeyPair(str(seat.order_id), '0'),
                cond_expr='attribute_exists(sk)',
                exc_cls=OrderNotFoundError,
                table_name=ORDER_TABLE,
            )
            transact.put(
                item={
                    'id': seat.order_id,
                    'sk': f'ENROLLMENT#{enrollment.id}',
                    'user': user.model_dump(),
                    'course': course.model_dump(),
                    'status': 'SCHEDULED',
                    'scheduled_at': now_,
                    'created_at': now_,
                },
                table_name=ORDER_TABLE,
            )
            transact.delete(
                key=KeyPair(
                    f'SEAT#ORG#{org.id}',
                    f'ORDER#{seat.order_id}#ENROLLMENT#{enrollment.id}',
                ),
                cond_expr='attribute_exists(sk)',
                exc_cls=SeatNotFoundError,
            )

        transact.put(
            item={
                'id': 'LOCK#SCHEDULED',
                'sk': lock_hash,
                'scheduled': {
                    'id': pk,
                    'sk': sk,
                },
                'ttl': ttl(start_dt=scheduled_for),
                'created_at': now_,
            },
            cond_expr='attribute_not_exists(sk)',
            exc_cls=DeduplicationConflictError,
        )

        if subscription:
            transact.condition(
                key=KeyPair('SUBSCRIPTION', f'ORG#{org.id}'),
                cond_expr='attribute_exists(sk)',
                exc_cls=SubscriptionRequiredError,
                table_name=USER_TABLE,
            )
            transact.condition(
                key=KeyPair(str(org.id), 'METADATA#SUBSCRIPTION'),
                cond_expr='billing_day = :billing_day',
                expr_attr_values={
                    ':billing_day': subscription.billing_day,
                },
                exc_cls=SubscriptionConflictError,
                table_name=USER_TABLE,
            )
            transact.condition(
                key=KeyPair('SUBSCRIPTION#FROZEN', f'ORG#{org.id}'),
                cond_expr='attribute_not_exists(sk)',
                exc_cls=SubscriptionFrozenError,
                table_name=USER_TABLE,
            )

    return enrollment


def _date_to_midnight(dt: date) -> datetime:
    return datetime.combine(dt, time(0, 0)).replace(tzinfo=pytz.timezone(TZ))