Amazon Locker

from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from threading import RLock
import hashlib
import secrets
import uuid


class LockerSize(Enum):
    SMALL = (30, 30, 10)
    MEDIUM = (45, 45, 30)
    LARGE = (60, 60, 45)
    XL = (90, 90, 60)

    def __init__(self, length: int, width: int, height: int):
        self.length = length
        self.width = width
        self.height = height


class LockerState(Enum):
    AVAILABLE = "available"
    RESERVED = "reserved"
    OCCUPIED = "occupied"
    OUT_OF_SERVICE = "out_of_service"


class ReservationStatus(Enum):
    ACTIVE = "active"
    COMPLETED = "completed"
    EXPIRED = "expired"


@dataclass(frozen=True)
class Dimensions:
    length: int
    width: int
    height: int

    def fits_inside(self, other: "Dimensions") -> bool:
        # Allow rotation: sort both sorted descending and compare coordinate-wise.
        a = sorted([self.length, self.width, self.height], reverse=True)
        b = sorted([other.length, other.width, other.height], reverse=True)
        return all(x <= y for x, y in zip(a, b))

    @property
    def volume(self) -> int:
        return self.length * self.width * self.height


@dataclass
class Package:
    id: str
    recipient: str
    dims: Dimensions
    weight_grams: int
    deadline: datetime


class Locker:
    def __init__(self, locker_id: str, size: LockerSize):
        self.id = locker_id
        self.size = size
        self.dims = Dimensions(*size.value)
        self.state = LockerState.AVAILABLE
        self.reservation_id: str | None = None

    @property
    def is_available(self) -> bool:
        return self.state == LockerState.AVAILABLE


class OtpService:
    """OTPs stored hashed; verification uses constant-time compare."""

    @staticmethod
    def generate() -> tuple[str, str]:
        # 6-digit code, returned alongside its hash. Station stores only the hash.
        code = f"{secrets.randbelow(1_000_000):06d}"
        return code, OtpService._hash(code)

    @staticmethod
    def verify(code: str, stored_hash: str) -> bool:
        return secrets.compare_digest(OtpService._hash(code), stored_hash)

    @staticmethod
    def _hash(code: str) -> str:
        # SHA-256 with a fixed salt here for brevity; production wants per-OTP salt or argon2.
        return hashlib.sha256(("locker-salt-" + code).encode()).hexdigest()


@dataclass
class Reservation:
    id: str
    package_id: str
    locker_id: str
    otp_hash: str
    expires_at: datetime
    status: ReservationStatus = ReservationStatus.ACTIVE


class AllocationStrategy(ABC):
    @abstractmethod
    def pick(self, lockers: list[Locker], package: Package) -> Locker | None: ...


class SmallestThatFits(AllocationStrategy):
    """Preserve larger lockers for larger packages."""

    def pick(self, lockers: list[Locker], package: Package) -> Locker | None:
        candidates = [
            l for l in lockers
            if l.is_available and package.dims.fits_inside(l.dims)
        ]
        if not candidates:
            return None
        return min(candidates, key=lambda l: l.dims.volume)


class LockerStation:
    def __init__(self, station_id: str, lockers: list[Locker],
                 strategy: AllocationStrategy | None = None):
        if not lockers:
            raise ValueError("station must have at least one locker")
        self.id = station_id
        self._lockers = {l.id: l for l in lockers}
        self._strategy = strategy or SmallestThatFits()
        self._reservations: dict[str, Reservation] = {}
        self._packages: dict[str, Package] = {}
        self._lock = RLock()

    def drop_off(self, package: Package, hold_ttl: timedelta = timedelta(days=3)) -> tuple[Reservation, str]:
        """Driver drops a package. Returns the reservation and the OTP (shown once to recipient).

        The reservation expires at the earlier of:
          - package.deadline (absolute cutoff the carrier committed to the sender)
          - now + hold_ttl (default station hold window)
        """
        with self._lock:
            now = datetime.utcnow()
            if package.deadline <= now:
                raise RuntimeError("Package deadline is already in the past")
            locker = self._strategy.pick(list(self._lockers.values()), package)
            if locker is None:
                raise RuntimeError("No locker fits this package at this station")
            otp, otp_hash = OtpService.generate()
            expires_at = min(package.deadline, now + hold_ttl)
            reservation = Reservation(
                id=str(uuid.uuid4())[:8],
                package_id=package.id,
                locker_id=locker.id,
                otp_hash=otp_hash,
                expires_at=expires_at,
            )
            locker.state = LockerState.OCCUPIED
            locker.reservation_id = reservation.id
            self._reservations[reservation.id] = reservation
            self._packages[package.id] = package
            return reservation, otp

    def pickup(self, reservation_id: str, otp: str) -> Package:
        with self._lock:
            reservation = self._reservations.get(reservation_id)
            # Unified error — don't leak existence.
            if reservation is None or reservation.status != ReservationStatus.ACTIVE:
                raise PermissionError("Invalid reservation or code")
            if datetime.utcnow() >= reservation.expires_at:
                self._expire_reservation(reservation)
                raise PermissionError("Invalid reservation or code")
            if not OtpService.verify(otp, reservation.otp_hash):
                raise PermissionError("Invalid reservation or code")

            locker = self._lockers[reservation.locker_id]
            pkg = self._packages.pop(reservation.package_id)
            reservation.status = ReservationStatus.COMPLETED
            locker.state = LockerState.AVAILABLE
            locker.reservation_id = None
            return pkg

    def sweep_expired(self) -> int:
        """Background hygiene — reclaim lockers whose pickup deadlines have passed."""
        with self._lock:
            now = datetime.utcnow()
            expired = [r for r in self._reservations.values()
                       if r.status == ReservationStatus.ACTIVE and now >= r.expires_at]
            for r in expired:
                self._expire_reservation(r)
            return len(expired)

    def _expire_reservation(self, r: Reservation) -> None:
        r.status = ReservationStatus.EXPIRED
        locker = self._lockers[r.locker_id]
        locker.state = LockerState.AVAILABLE
        locker.reservation_id = None
        # In production this triggers a "return to sender" workflow.

    def availability_by_size(self) -> dict[LockerSize, int]:
        with self._lock:
            counts: dict[LockerSize, int] = {}
            for locker in self._lockers.values():
                if locker.is_available:
                    counts[locker.size] = counts.get(locker.size, 0) + 1
            return counts


if __name__ == "__main__":
    # Station with 2 lockers per size.
    lockers = []
    for i, size in enumerate(LockerSize):
        for j in range(2):
            lockers.append(Locker(f"L-{size.name}-{j}", size))
    station = LockerStation("AMZL-001", lockers)

    print("Initial availability:", {s.name: c for s, c in station.availability_by_size().items()})

    # Small package fits in SMALL — should go there first.
    small_pkg = Package(
        id="P-001", recipient="alice@example.com",
        dims=Dimensions(20, 20, 5), weight_grams=500,
        deadline=datetime.utcnow() + timedelta(days=2),
    )
    res1, otp1 = station.drop_off(small_pkg)
    assigned_size = next(l.size for l in lockers if l.id == res1.locker_id)
    print(f"Small package -> {assigned_size.name} locker (code: {otp1})")

    # Too-big package that's still valid for LARGE.
    big_pkg = Package(
        id="P-002", recipient="bob@example.com",
        dims=Dimensions(55, 55, 40), weight_grams=8000,
        deadline=datetime.utcnow() + timedelta(days=2),
    )
    res2, otp2 = station.drop_off(big_pkg)
    print(f"Big package -> {next(l.size.name for l in lockers if l.id == res2.locker_id)}")

    # Wrong OTP → rejected.
    try:
        station.pickup(res1.id, "000000")
    except PermissionError as e:
        print(f"Wrong OTP rejected: {e}")

    # Correct OTP → returned package.
    retrieved = station.pickup(res1.id, otp1)
    print(f"Picked up: {retrieved.id}")
    print("After pickup:", {s.name: c for s, c in station.availability_by_size().items()})

    # Package too large for any locker.
    huge_pkg = Package(
        id="P-FRIDGE", recipient="carol@example.com",
        dims=Dimensions(200, 80, 80), weight_grams=50000,
        deadline=datetime.utcnow() + timedelta(days=2),
    )
    try:
        station.drop_off(huge_pkg)
    except RuntimeError as e:
        print(f"Too-big package rejected: {e}")

    # Package deadline earlier than the default 3-day hold caps the reservation.
    tight_pkg = Package(
        id="P-TIGHT", recipient="dan@example.com",
        dims=Dimensions(15, 15, 5), weight_grams=200,
        deadline=datetime.utcnow() + timedelta(hours=6),
    )
    res_tight, _ = station.drop_off(tight_pkg, hold_ttl=timedelta(days=3))
    hours_to_expiry = (res_tight.expires_at - datetime.utcnow()).total_seconds() / 3600
    print(f"Tight-deadline reservation expires in ~{hours_to_expiry:.1f}h (capped by package.deadline)")
    assert hours_to_expiry < 7