Blinkit Quick Commerce.

# serviceability/store_selector.py

"""
Picks the serving dark store for a customer location.

The static part (which stores could ever serve a cell) is precomputed.
The live part (is that store open for this cell, how busy is it) comes
from the promise engine. Nothing here touches the stock database.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set

import h3  # h3-py v4

RES = 9                 # ~0.105 km² cells, ~200 m edges
PICK_PACK_MIN = 2.5     # median pick + pack, refreshed per store in practice
MAX_RIDE_MIN = 20.0     # fallback when a (store, cell) has no ride history


@dataclass
class StoreLive:
    store_id: str
    mode: str                       # OPEN | STRETCHED | THROTTLED | CLOSED
    backlog_min: float              # minutes of picking work already queued
    open_cells: Set[str]            # cells it accepts right now (shrinks under load)
    ride_p90_min: Dict[str, float] = field(default_factory=dict)  # cell -> p90 ride


@dataclass
class Choice:
    store_id: str
    cell: str
    eta_min: float
    coverage: float                 # share of cart lines in stock at this store


class StoreSelector:
    def __init__(self, cell_index: Dict[str, List[str]], live: Dict[str, StoreLive], stock):
        self.cell_index = cell_index    # cell -> candidate store ids (static)
        self.live = live                # store id -> StoreLive (refreshed every few s)
        self.stock = stock              # availability cache: in_stock(store, sku) -> bool

    def choose(self, lat: float, lng: float, cart_skus: List[str]) -> Optional[Choice]:
        cell = h3.latlng_to_cell(lat, lng, RES)
        best: Optional[Choice] = None

        for store_id in self.cell_index.get(cell, []):
            st = self.live.get(store_id)
            if st is None or st.mode == "CLOSED" or cell not in st.open_cells:
                continue

            in_stock = sum(1 for sku in cart_skus if self.stock.in_stock(store_id, sku))
            coverage = in_stock / len(cart_skus) if cart_skus else 1.0
            eta = st.backlog_min + PICK_PACK_MIN + st.ride_p90_min.get(cell, MAX_RIDE_MIN)

            cand = Choice(store_id, cell, round(eta, 1), coverage)
            # Coverage first (a missing item is worse than a minute), then ETA.
            if best is None or (cand.coverage, -cand.eta_min) > (best.coverage, -best.eta_min):
                best = cand

        return best  # None means "not serviceable right now"

# inventory/reservations.py

"""
Stock reservations for one dark store, in an ordinary SQL database.

Per store, reservations arrive at tens a minute, so a plain transaction
with a conditional UPDATE is enough: no distributed lock, no counter
cache. The tables are sharded by store_id; every statement below
touches one store, so it runs on one shard.

stock(store_id, sku_id, on_hand, reserved, buffer)      PK (store_id, sku_id)
reservations(order_id PK, store_id, status, expires_at)
reservation_lines(order_id, sku_id, qty)
"""

from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Dict, List, Optional

HOLD_TTL = timedelta(minutes=5)


@dataclass
class Shortfall:
    sku_id: str
    wanted: int
    available: int


class StoreInventory:
    def __init__(self, db):
        self.db = db

    def reserve(self, order_id: str, store_id: str, lines: Dict[str, int]) -> List[Shortfall]:
        """
        All-or-nothing reservation for one basket. Returns [] on success,
        or the lines that are short (so the app can say "only 1 left").
        Idempotent on order_id: a retried checkout returns the same result.
        """
        with self.db.transaction() as tx:
            inserted = tx.execute(
                "INSERT INTO reservations (order_id, store_id, status, expires_at) "
                "VALUES (%s, %s, 'HELD', %s) ON CONFLICT (order_id) DO NOTHING",
                (order_id, store_id, datetime.utcnow() + HOLD_TTL),
            ).rowcount
            if inserted == 0:
                status = tx.scalar("SELECT status FROM reservations WHERE order_id = %s", (order_id,))
                if status in ("HELD", "COMMITTED"):
                    return []           # same checkout retried: already reserved
                # an earlier attempt with this id failed or expired: start a new hold
                tx.execute(
                    "UPDATE reservations SET status = 'HELD', expires_at = %s WHERE order_id = %s",
                    (datetime.utcnow() + HOLD_TTL, order_id),
                )
                tx.execute("DELETE FROM reservation_lines WHERE order_id = %s", (order_id,))

            short: List[Shortfall] = []
            # Fixed order (by sku_id) so two baskets never lock the same rows in opposite orders.
            for sku_id in sorted(lines):
                qty = lines[sku_id]
                ok = tx.execute(
                    "UPDATE stock SET reserved = reserved + %s "
                    "WHERE store_id = %s AND sku_id = %s "
                    "AND on_hand - reserved - buffer >= %s",
                    (qty, store_id, sku_id, qty),
                ).rowcount
                if ok == 0:
                    avail = tx.scalar(
                        "SELECT GREATEST(on_hand - reserved - buffer, 0) FROM stock "
                        "WHERE store_id = %s AND sku_id = %s",
                        (store_id, sku_id),
                    ) or 0
                    short.append(Shortfall(sku_id, qty, avail))
                else:
                    tx.execute(
                        "INSERT INTO reservation_lines (order_id, sku_id, qty) VALUES (%s, %s, %s)",
                        (order_id, sku_id, qty),
                    )

            if short:
                tx.rollback()           # nothing held: the customer edits the cart and retries
                return short
            return []

    def extend_hold(self, order_id: str, by: timedelta) -> None:
        """Payment outcome unknown: keep the stock while the reconciler asks the gateway."""
        with self.db.transaction() as tx:
            tx.execute(
                "UPDATE reservations SET expires_at = expires_at + %s "
                "WHERE order_id = %s AND status = 'HELD'",
                (by, order_id),
            )

    def expire_holds(self, now: datetime) -> None:
        """The sweeper, every few seconds: release holds whose time is up."""
        for (order_id,) in self.db.rows(
            "SELECT order_id FROM reservations WHERE status = 'HELD' AND expires_at < %s", (now,)
        ):
            self.release(order_id)

    def commit(self, order_id: str) -> bool:
        """Payment succeeded: the hold stops expiring. False if it already expired."""
        with self.db.transaction() as tx:
            return tx.execute(
                "UPDATE reservations SET status = 'COMMITTED', expires_at = NULL "
                "WHERE order_id = %s AND status = 'HELD'",
                (order_id,),
            ).rowcount == 1

    def release(self, order_id: str) -> None:
        """Payment failed or the hold expired. Safe to call twice."""
        with self.db.transaction() as tx:
            # Only the call that moves HELD -> RELEASED gives the stock back.
            moved = tx.execute(
                "UPDATE reservations SET status = 'RELEASED' "
                "WHERE order_id = %s AND status = 'HELD' RETURNING store_id",
                (order_id,),
            ).fetchone()
            if moved is None:
                return
            store_id = moved[0]
            for sku_id, qty in tx.rows(
                "SELECT sku_id, qty FROM reservation_lines WHERE order_id = %s ORDER BY sku_id",
                (order_id,),
            ):
                tx.execute(
                    "UPDATE stock SET reserved = reserved - %s WHERE store_id = %s AND sku_id = %s",
                    (qty, store_id, sku_id),
                )

    def scan_out(self, store_id: str, order_id: str, sku_id: str, qty: int) -> None:
        """The picker scanned the item into the bag: it leaves the shelf."""
        with self.db.transaction() as tx:
            tx.execute(
                "UPDATE stock SET on_hand = on_hand - %s, reserved = reserved - %s "
                "WHERE store_id = %s AND sku_id = %s",
                (qty, qty, store_id, sku_id),
            )

    def not_found(self, store_id: str, order_id: str, sku_id: str, qty: int, counted: int) -> None:
        """
        The bin held fewer than the system thought. (A partial pick is a
        scan_out for what was there, then not_found for the rest.) Trust the
        picker: `counted` is what is physically left, usually 0. Release this
        line's hold and queue a recount. This sketch keeps one bin per SKU.
        """
        with self.db.transaction() as tx:
            tx.execute(
                "UPDATE stock SET reserved = reserved - %s, on_hand = LEAST(on_hand, %s) "
                "WHERE store_id = %s AND sku_id = %s",
                (qty, counted, store_id, sku_id),
            )
            tx.execute(
                "INSERT INTO recount_tasks (store_id, sku_id, reason) VALUES (%s, %s, 'NOT_FOUND')",
                (store_id, sku_id),
            )

# orders/saga.py

"""
Order saga for one checkout. Each step is idempotent, so the whole
place_order call can be retried safely with the same key.
"""

from dataclasses import dataclass
from datetime import timedelta
from typing import Dict, Optional


@dataclass
class Order:
    order_id: str
    idem_key: str
    store_id: str
    lines: Dict[str, int]
    state: str
    payment_ref: Optional[str] = None


class OrderSaga:
    def __init__(self, orders, inventory, payments, outbox, scheduler):
        self.orders = orders            # orders DB (same shard as the store's stock)
        self.inv = inventory            # StoreInventory from Deep Dive 2
        self.pay = payments
        self.outbox = outbox
        self.scheduler = scheduler

    def place_order(self, idem_key: str, store_id: str, lines: Dict[str, int], method: str) -> Order:
        order = self.orders.get_or_create(idem_key, store_id, lines)   # unique on idem_key
        if order.state not in ("CREATED", "RESERVED"):
            return order                # already decided: a retry gets the same answer

        if order.state == "CREATED":
            short = self.inv.reserve(order.order_id, store_id, lines)
            if short:
                return self.orders.transition(order, "REJECTED", shortfall=short)
            order = self.orders.transition(order, "RESERVED")

        if method == "COD":
            return self._confirm(order)

        result = self.pay.charge(reference=order.order_id, idempotency_key=idem_key)
        if result.status == "SUCCESS":
            order.payment_ref = result.reference
            return self._confirm(order)
        if result.status == "TIMEOUT":
            # Don't guess. Keep the stock a little longer and ask the gateway.
            self.inv.extend_hold(order.order_id, timedelta(minutes=10))
            self.scheduler.reconcile_payment(order.order_id, after=timedelta(seconds=30))
            return self.orders.transition(order, "PAYMENT_PENDING")

        self.inv.release(order.order_id)
        return self.orders.transition(order, "CANCELLED", reason="PAYMENT_FAILED")

    def _confirm(self, order: Order) -> Order:
        if not self.inv.commit(order.order_id):
            # The hold expired while the customer was paying. Try once more.
            if self.inv.reserve(order.order_id, order.store_id, order.lines) or not self.inv.commit(order.order_id):
                if order.payment_ref:
                    self.pay.refund(order.payment_ref, idempotency_key=f"refund:{order.order_id}")
                return self.orders.transition(order, "REFUNDED", reason="STOCK_GONE")
        # The state change and its event commit together (transactional outbox).
        with self.orders.transaction() as tx:
            order = self.orders.transition(order, "CONFIRMED", tx=tx)
            self.outbox.add(tx, "order.confirmed", {"order_id": order.order_id, "store_id": order.store_id})
        return order

# store/picking.py

"""
Picker queue and pick lists for one dark store.
"""

import heapq
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

ZONE_ORDER = {"AMBIENT": 0, "CHILLED": 1, "FROZEN": 2}   # cold items last
SECONDS_PER_LINE = 15     # walk + reach + scan, measured per store in practice
PACK_SECONDS = 40


@dataclass(order=True)
class QueuedOrder:
    pack_by: float                                  # epoch seconds; heap key
    order_id: str = field(compare=False)
    lines: List[Tuple[str, int]] = field(compare=False)   # (sku_id, qty)


class PickerQueue:
    def __init__(self, bins: Dict[str, Tuple[str, int]]):
        self.bins = bins            # sku_id -> (zone, walk_seq)
        self.heap: List[QueuedOrder] = []

    def push(self, order: QueuedOrder) -> None:
        heapq.heappush(self.heap, order)

    def pick_list(self, lines: List[Tuple[str, int]]) -> List[Tuple[str, int]]:
        """One walk through the store: zone order first, then the bins' walk sequence."""
        return sorted(lines, key=lambda l: (ZONE_ORDER[self.bins[l[0]][0]], self.bins[l[0]][1]))

    def est_seconds(self, lines: List[Tuple[str, int]]) -> float:
        return len(lines) * SECONDS_PER_LINE + PACK_SECONDS

    def next_task(self, now: float) -> Optional[List[QueuedOrder]]:
        """
        Give a free picker the earliest-deadline order, plus a second
        order on the same trolley if both still finish in time.
        """
        if not self.heap:
            return None
        first = heapq.heappop(self.heap)
        if self.heap and len(first.lines) <= 8 and len(self.heap[0].lines) <= 8:
            second = self.heap[0]
            zones_a = {self.bins[s][0] for s, _ in first.lines}
            zones_b = {self.bins[s][0] for s, _ in second.lines}
            # The shared walk takes the lines of both; the second bag's packing adds its own time.
            both = self.est_seconds(first.lines + second.lines) + PACK_SECONDS
            if zones_b <= zones_a and now + both <= first.pack_by and now + both <= second.pack_by:
                heapq.heappop(self.heap)
                return [first, second]
        return [first]

# dispatch/batching.py

"""
Decides whether two ready orders share one rider. Ride times come
from a model trained on the store's own past trips (by hour, weather
and destination cell), as p90 minutes, so "on time" is judged
conservatively.
"""

from dataclasses import dataclass
from typing import List, Optional, Tuple

MAX_DETOUR_MIN = 2.0


@dataclass
class ReadyOrder:
    order_id: str
    cell: str                   # H3 cell of the drop
    ready_at_min: float         # minutes from now until packed (0 if on the rack)
    promise_at_min: float       # minutes from now until the shown arrival time


class Batcher:
    def __init__(self, ride_model, handover_min: float = 1.0):
        self.ride = ride_model      # ride.p90(from_cell_or_store, to_cell) -> minutes
        self.handover = handover_min

    def plan(self, a: ReadyOrder, b: ReadyOrder, rider_free_in_min: float) -> Optional[Tuple[ReadyOrder, ReadyOrder]]:
        """Return the (first, second) drop order if batching keeps both promises."""
        start = max(a.ready_at_min, b.ready_at_min, rider_free_in_min) + self.handover
        best = None
        for first, second in ((a, b), (b, a)):
            t_first = start + self.ride.p90("STORE", first.cell)
            t_second = t_first + self.handover + self.ride.p90(first.cell, second.cell)
            solo_second = max(second.ready_at_min, rider_free_in_min) + self.handover + self.ride.p90("STORE", second.cell)
            detour = t_second - solo_second
            if t_first <= first.promise_at_min and t_second <= second.promise_at_min and detour <= MAX_DETOUR_MIN:
                if best is None or t_second < best[0]:
                    best = (t_second, (first, second))
        return best[1] if best else None


def pick_rider(riders: List[dict], needed_at_min: float) -> Optional[dict]:
    """Longest-waiting rider who will be at the store in time (fairness first)."""
    eligible = [r for r in riders if r["back_in_min"] <= needed_at_min]
    return max(eligible, key=lambda r: r["waiting_since_min"], default=None)

# promise/engine.py

"""
Per-store promise engine. Runs every few seconds per store and
publishes (mode, open cells, per-cell estimates) to serviceability.
"""

import math
from dataclasses import dataclass
from typing import Dict, List

PICK_P90_MIN = 3.5          # pick + pack p90 for a typical basket at this store
HANDOVER_MIN = 0.5


@dataclass
class StoreNow:
    queued_orders: int
    pickers_active: int
    riders_at_store: int
    riders_back_in_min: List[float]     # predicted return times of riders out
    promise_kept_last_hour: float       # share delivered by the shown time
    healthy: bool                       # power, network, handhelds


class PromiseEngine:
    # (enter threshold, leave threshold) for backlog minutes, per mode
    ENTER = {"STRETCHED": 4.0, "THROTTLED": 8.0, "CLOSED": 15.0}
    LEAVE = {"STRETCHED": 2.0, "THROTTLED": 5.0, "CLOSED": 8.0}
    ORDER = ["OPEN", "STRETCHED", "THROTTLED", "CLOSED"]

    def __init__(self, ride_model, outer_cells: List[str], all_cells: List[str]):
        self.ride = ride_model
        self.outer = set(outer_cells)       # cells with the longest rides
        self.cells = all_cells
        self.mode = "OPEN"
        self.calm_since = None              # when the backlog last fell below LEAVE

    def backlog_min(self, s: StoreNow) -> float:
        if s.pickers_active == 0:
            return math.inf
        return s.queued_orders / s.pickers_active * 2.5

    def rider_wait_min(self, s: StoreNow) -> float:
        if s.riders_at_store > 0:
            return 0.0
        return min(s.riders_back_in_min, default=15.0)

    def estimate(self, s: StoreNow, cell: str) -> int:
        eta = (self.backlog_min(s) + PICK_P90_MIN
               + max(0.0, self.rider_wait_min(s) - PICK_P90_MIN)   # rider wait overlaps picking
               + HANDOVER_MIN + self.ride.p90("STORE", cell))
        return int(math.ceil(eta))

    def pressure(self, s: StoreNow, nxt: str) -> bool:
        """Should the store move up to mode `nxt`?"""
        if self.backlog_min(s) > self.ENTER[nxt]:
            return True
        if nxt == "STRETCHED" and self.rider_wait_min(s) > PICK_P90_MIN:
            return True     # bags would be packed before any rider is back
        if nxt == "THROTTLED" and s.promise_kept_last_hour < 0.85:
            return True     # the model is wrong about tonight
        return False

    def step(self, s: StoreNow, now_min: float) -> Dict:
        b = self.backlog_min(s)
        i = self.ORDER.index(self.mode)

        if not s.healthy:
            self.mode = "CLOSED"
            self.calm_since = None
        elif i < 3 and self.pressure(s, self.ORDER[i + 1]):
            self.mode = self.ORDER[i + 1]           # tighten at once, one step at a time
            self.calm_since = None
        elif i > 0 and b < self.LEAVE[self.mode]:
            self.calm_since = self.calm_since if self.calm_since is not None else now_min
            hold = 5.0 if self.mode == "CLOSED" else 3.0
            if now_min - self.calm_since >= hold:
                self.mode = self.ORDER[i - 1]       # relax slowly
                self.calm_since = None
        else:
            self.calm_since = None

        open_cells = [] if self.mode == "CLOSED" else (
            [c for c in self.cells if c not in self.outer] if self.mode == "THROTTLED" else list(self.cells)
        )
        return {
            "mode": self.mode,
            "open_cells": open_cells,
            "estimates": {c: self.estimate(s, c) for c in open_cells},
        }