undercooked/tests/sim/helpers.py

"""The headless full-round harness (phase 04 task 06) — test code only.

`src/` never imports this module. It gives `test_full_rounds.py`'s scripted
full rounds and `test_perf_full_pipeline.py`'s budget run:

* `make_level_1_state` — a tick-0 round on the shipped level 1
  (`levels/01_first_kitchen.yaml`) with the shipped recipe book, built
  through the real seam (`content.load_level_file`,
  `state_build.build_state`) and then seeded the way a phase-05 room would
  seed it: the spawn shelves' stock from `level.station_items` (the room's
  job — `sim/stations.py` documents the stand-in) and the prop stock
  counters from the level's `station_limits` (the wash-cycle test passes
  `plates=1`). `target` overrides `target_score` for the two-player and
  wash-cycle tests, whose bar a single first dish (≈ 58 points) cannot
  clear; the shipped 45 is the default.

* The scripted cook: a BFS pathfinder over the walkable grid (a non-
  `PASS_THROUGH_STATIONS` fixture blocks its cell, the walkable trio never
  does) and the one-button controller the action gate prescribes — a
  release tick between presses (one press, one effect), an explicit face
  tick before the press that uses it (the gate targets with the
  pre-movement facing), and `action=True` held while a station's work runs
  (the hold is the progress driver, `sim/stations.py` stage 4).

* `decide`, the reactive "serve everything possible" policy: a free hand
  picks ready outputs, picks a tabled object, washes a dirty work fixture
  in place (the stuck prop's fix) and picks stock from the cook's own
  shelf first; a carried raw goes to the chop board (the set-down starts
  the chop), a chopped to the stove (the set-down starts the cook), a
  cooked to the prep (the press plates it), a plated to the window (waiting
  there while the rail is empty — a plated dish does not decay in the
  hand). A full hand at a dirty fixture tables the item on the chop board
  (the board works raw only, so the set-down stays a marked object — the
  hand-off's grabbable), frees the hand for the in-place wash, and picks
  it back. A fixture working for the other cook is waited on, not pressed.

* `run_round` — drive `max_ticks` substeps of the whole `step()` with one
  frame per cook per tick, returning the per-tick states, the per-tick
  event lists and the tick the verdict first landed (if it did).
"""

from __future__ import annotations

import math
from collections import deque
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field, replace
from pathlib import Path

from undercooked.sim.collision import PASS_THROUGH_STATIONS
from undercooked.sim.content import load_level_file
from undercooked.sim.model import (
    WALKABLE,
    Facing,
    Item,
    ItemForm,
    ItemKind,
    KitchenState,
    Player,
    Station,
    StationKind,
    StationState,
    Vec2,
)
from undercooked.sim.movement import RADIUS
from undercooked.sim.recipe import RecipeBook, load_recipe_book
from undercooked.sim.state_build import SeatSpec, build_state
from undercooked.sim.step import STEP_DT, PlayerInput, step

REPO_ROOT = Path(__file__).resolve().parents[2]
LEVEL1_YAML = REPO_ROOT / "levels" / "01_first_kitchen.yaml"
RECIPES_YAML = REPO_ROOT / "levels" / "recipes.yaml"

__all__ = [
    "LEVEL1_YAML",
    "RECIPES_YAML",
    "CookController",
    "bfs_path",
    "decide",
    "load_book",
    "make_level_1_state",
    "run_round",
    "shelf_id_for",
    "stand_options",
    "walkable_cells",
]

#: The arrival band around a stand cell's centre: wide enough to absorb the
#: deceleration drift and a soft-push nudge, narrow enough that the press
#: stays inside `REACH + FORGIVE` (1.0 + the band is well under 1.6).
GOTO_BAND = 0.3
#: The stall test's slack against the low-watermark distance: a stall is
#: "worse than the best so far by this much" (a soft-push jitter swings
#: the distance by ~0.15 per tick, so a naive "did it get closer" counter
#: resets every other tick forever — the low watermark does not).
GOTO_STALL_SLACK = 0.05
#: The reset bar against the low-watermark distance: only beating the best
#: so far by a full step is genuine progress. A waypoint hand-off jumps
#: the distance by up to ~1.0 and takes ~7 ticks to walk back down; with
#: this bar the counter creeps to at most ~7 per hand-off and resets
#: (below `GOTO_STUCK_TICKS`) before the next one, while a real block
#: (which never beats the watermark) accumulates past the brake.
GOTO_PROGRESS = 0.08
#: Stalled-approach ticks before the detour (or, when no detour exists,
#: the brake that lets the way clear).
GOTO_STUCK_TICKS = 15


# --- the round builder ---------------------------------------------------------


def load_book() -> RecipeBook:
    """The shipped recipe book (the levels' `recipes.yaml`)."""
    return load_recipe_book(RECIPES_YAML)


def make_level_1_state(
    seed: int = 0x5EED_0406,
    seats: int = 1,
    target: int | None = None,
    plates: int = 99,
    pots: int = 99,
    shelf_stock: int = 8,
) -> KitchenState:
    """Tick 0 of a level-1 round, built and seeded like a phase-05 room.

    Args:
        seed: the round seed (the `rng_state`'s start).
        seats: how many cooks (0 for the all-idle rounds).
        target: `target_score` override; the shipped value when `None`.
        plates / pots: the clean prop stock (the level's `station_limits`
            stand-in; the wash-cycle test's 1-plate kitchen).
        shelf_stock: items seeded on every spawn shelf.
    """
    level = load_level_file(LEVEL1_YAML)
    book = load_book()
    seat_specs = [
        SeatSpec(player_id=f"p{index}", seat=index, display_name=f"cook {index + 1}")
        for index in range(seats)
    ]
    state = build_state(level, seat_specs, seed, book)
    # The test's stand-in for the phase-05 room's seeding (see the module
    # docstring): shelf stock from the level's legend, the stock counters
    # from the level's limits.
    stations = list(state.stations)
    for index, station in enumerate(stations):
        if station.kind is StationKind.SPAWN_SHELF:
            kind = level.station_items[station.id]
            stations[index] = replace(
                station, contents=(Item(kind, ItemForm.RAW),) * shelf_stock
            )
    updated = replace(
        state,
        stations=tuple(stations),
        counters={
            "plates_clean": plates,
            "plates_dirty": 0,
            "pots_clean": pots,
            "pots_dirty": 0,
        },
    )
    if target is not None:
        updated = replace(updated, target_score=target)
    return updated


def shelf_id_for(state: KitchenState, kind: ItemKind) -> str:
    """The id of the spawn shelf stocked with `kind` (a seeded one)."""
    for station in state.stations:
        if (
            station.kind is StationKind.SPAWN_SHELF
            and station.contents
            and station.contents[0].kind is kind
        ):
            return station.id
    msg = f"no {kind.value} shelf with stock"
    raise ValueError(msg)


# --- geometry -------------------------------------------------------------------


def _cell_of(pos: Vec2) -> tuple[int, int]:
    return math.floor(pos.x), math.floor(pos.y)


def _body_cells(pos: Vec2) -> frozenset[tuple[int, int]]:
    """The cells a body of radius `RADIUS` overlaps — the same footprint
    arithmetic as the collision (the eps-inset keeps a flush edge out of
    the cell it only touches), so the avoid set matches what actually
    blocks the other cook's body (a body straddles up to 4 cells)."""
    x0, x1 = math.floor(pos.x - RADIUS + 1e-6), math.floor(pos.x + RADIUS - 1e-6)
    y0, y1 = math.floor(pos.y - RADIUS + 1e-6), math.floor(pos.y + RADIUS - 1e-6)
    return frozenset(
        (x, y) for x in range(x0, x1 + 1) for y in range(y0, y1 + 1)
    )


def _station_cell(station: Station) -> tuple[int, int]:
    return int(station.x), int(station.y)


def walkable_cells(state: KitchenState) -> frozenset[tuple[int, int]]:
    """The cells a cook's body may occupy.

    Walkable terrain minus the cells of the fixtures that are not
    pass-through (`collision.PASS_THROUGH_STATIONS` never block — a cook
    stands *on* a serve window, a sink or a trash can).
    """
    blocked = {
        _station_cell(station)
        for station in state.stations
        if station.kind not in PASS_THROUGH_STATIONS
    }
    width = state.grid.width
    cells: set[tuple[int, int]] = set()
    for index, tile in enumerate(state.grid.tiles):
        if tile in WALKABLE:
            cell = (index % width, index // width)
            if cell not in blocked:
                cells.add(cell)
    return frozenset(cells)


def stand_options(
    state: KitchenState, station: Station
) -> list[tuple[tuple[int, int], Facing]]:
    """`(stand cell, facing)` for every walkable cell orthogonally adjacent
    to the station's cell — the cells where a press reaches the fixture
    (centre-to-centre 1.0, the reach rule's `REACH`)."""
    cells = walkable_cells(state)
    sx, sy = _station_cell(station)
    # (offset from the station to the stand, the facing *toward* the station)
    options: list[tuple[tuple[int, int], Facing]] = []
    for dx, dy, facing in (
        (1, 0, Facing.LEFT),
        (-1, 0, Facing.RIGHT),
        (0, -1, Facing.DOWN),
        (0, 1, Facing.UP),
    ):
        cell = (sx + dx, sy + dy)
        if cell in cells:
            options.append((cell, facing))
    return options


def bfs_path(
    state: KitchenState,
    start: tuple[int, int],
    goal: tuple[int, int],
    avoid: frozenset[tuple[int, int]] | None = None,
) -> list[tuple[int, int]]:
    """BFS over the walkable cells from `start` to `goal` (excluding `start`).

    `avoid` cells count as blocked (the other cooks' cells, for the corridor
    detour). The start may be off-grid-legal (a cook standing on a
    pass-through station cell is still walkable; a hand-placed one is not)
    — it snaps to the nearest walkable cell. `ValueError` when the goal is
    unreachable.
    """
    cells = walkable_cells(state)
    if avoid:
        cells = cells - set(avoid)
    origin = start
    if start not in cells:
        if not cells:
            msg = f"no walkable cells left from {start}"
            raise ValueError(msg)
        origin = min(cells, key=lambda cell: abs(cell[0] - start[0]) + abs(cell[1] - start[1]))
    if goal not in cells:
        msg = f"goal cell {goal} is not walkable"
        raise ValueError(msg)
    previous: dict[tuple[int, int], tuple[int, int] | None] = {origin: None}
    queue = deque([origin])
    while queue:
        current = queue.popleft()
        if current == goal:
            break
        for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            neighbour = (current[0] + dx, current[1] + dy)
            if neighbour in cells and neighbour not in previous:
                previous[neighbour] = current
                queue.append(neighbour)
    if goal not in previous:
        msg = f"no path from {start} to {goal}"
        raise ValueError(msg)
    path: list[tuple[int, int]] = []
    current: tuple[int, int] | None = goal
    while current is not None:
        path.append(current)
        current = previous[current]
    path.reverse()
    return path[1:]


def _station_by_id(state: KitchenState, station_id: str) -> Station:
    found = next((s for s in state.stations if s.id == station_id), None)
    if found is None:
        msg = f"station {station_id!r} is not in the state"
        raise ValueError(msg)
    return found


def _is_dirty(station: Station) -> bool:
    return station.dirty or station.state is StationState.DIRTY


def _face_to(player: Player, station: Station) -> Facing:
    """The facing that puts `station` in front of `player` (underfoot: any)."""
    px, py = _cell_of(player.pos)
    sx, sy = _station_cell(station)
    dx, dy = sx - px, sy - py
    if dx == 0 and dy == 0:
        return player.facing  # standing on a pass-through: underfoot reach
    if abs(dx) >= abs(dy):
        return Facing.RIGHT if dx > 0 else Facing.LEFT
    return Facing.DOWN if dy > 0 else Facing.UP


def _frame(
    player_id: str,
    tick: int,
    move: Vec2 | None = None,
    face: Facing | None = None,
    action: bool = False,
) -> PlayerInput:
    """One frame, stamped like a live room: `received_at_tick` is the tick the
    frame arrives on, `seq` a constant (one frame per cook per tick)."""
    return PlayerInput(
        player_id=player_id,
        move=move if move is not None else Vec2(),
        action=action,
        face=face,
        received_at_tick=tick,
    )


# --- the policy -------------------------------------------------------------------


def _station_free_predicate(station_id: str, me_id: str) -> Callable[[KitchenState], bool]:
    def free(state: KitchenState) -> bool:
        station = _station_by_id(state, station_id)
        return station.state is not StationState.WORKING or station.worker_id == me_id

    return free


def _board_free_predicate(station_id: str) -> Callable[[KitchenState], bool]:
    """The fixture is idle and empty — a set-down there is not refused
    (a set-down on a working or a loaded surface is a refusal, and a
    policy that walked in to refuse is the loop every full-hand branch
    used to hit)."""

    def free(state: KitchenState) -> bool:
        station = _station_by_id(state, station_id)
        return station.state is not StationState.WORKING and not station.contents

    return free


@dataclass
class _WaitUntil:
    """Idle frames until `predicate` (or the timeout — the replan re-issues)."""

    predicate: Callable[[KitchenState], bool]
    max_ticks: int
    left: int = field(init=False)

    def __post_init__(self) -> None:
        self.left = self.max_ticks

    def tick(
        self, state: KitchenState, player: Player, tick: int
    ) -> tuple[PlayerInput, bool]:
        if self.predicate(state) or self.left <= 0:
            return _frame(player.id, tick, face=player.facing), True
        self.left -= 1
        return _frame(player.id, tick, face=player.facing), False


@dataclass
class _Goto:
    """Walk the cook to a stand cell of the station and face the fixture.

    When the approach stalls — a head-on collision with the other cook in
    a one-tile corridor, which the soft push cannot resolve (the two 0.7-
    wide bodies grind against each other and a carried item would burn
    before either gets through) — the cook detours around the blockers'
    cells (BFS with them avoided) and drops the detour the moment the
    plain route is clear again. With no detour, the cook brakes instead.
    """

    index: int
    station_id: str
    cell: tuple[int, int]
    facing: Facing
    path: list[tuple[int, int]] = field(default_factory=list)
    stuck: int = 0
    low: float | None = None
    entered_hold: int = 0
    avoid: frozenset[tuple[int, int]] = field(default_factory=frozenset)

    def __init__(self, state: KitchenState, index: int, station_id: str) -> None:
        self.index = index
        self.station_id = station_id
        station = _station_by_id(state, station_id)
        options = stand_options(state, station)
        if not options:
            msg = f"station {station_id} has no walkable stand cell"
            raise ValueError(msg)
        start = _cell_of(state.players[index].pos)

        def distance_to(option: tuple[tuple[int, int], Facing]) -> int:
            return len(bfs_path(state, start, option[0]))

        self.cell, self.facing = min(
            options, key=lambda option: (distance_to(option), option[1].value)
        )
        self.path = bfs_path(state, start, self.cell)
        self.stuck = 0
        self.low = None
        self.entered_hold = 0
        self.avoid = frozenset()

    def tick(
        self, state: KitchenState, player: Player, tick: int
    ) -> tuple[PlayerInput, bool]:
        if self.already_there(player):
            self.avoid = frozenset()
            self.low = None
            return self._arrival_frame(player, tick)
        # The path from the *current* cell: BFS excludes the start, so the
        # first cell is always ahead (a cook on a waypoint walks to the next
        # one, not back at itself — that is how the old target oscillation
        # went away). A pushed-off cook re-roots the path on its own cell.
        # A cook mid-detour re-roots on the avoided path and returns to the
        # plain path the moment the blockers' cells are off it.
        here = _cell_of(player.pos)
        # Entering the previous first waypoint is one cell of progress —
        # debounced: a cook that sits in it for 3 ticks is walking (a
        # hand-off is legal, the distance jumped by up to ~1.0), while a
        # two-cell oscillation on a soft push spends one tick per cell and
        # never resets (the stall counter keeps accumulating). A cook
        # stalled in one cell is never in the entered cell, so it too
        # accumulates.
        entered = self.path[0] if self.path else None
        if entered is not None and here == entered:
            self.entered_hold = min(self.entered_hold + 1, 3)
        else:
            self.entered_hold = 0
        if self.entered_hold == 3:
            self.stuck = 0
        if self.avoid:
            plain = bfs_path(state, here, self.cell)
            # A co-located blocker also blocks the first step, which the
            # plain path's origin exclusion would not see: avoid the first
            # waypoint too, so the detour leaves in a free direction.
            extra = {plain[0]} if plain and here in self.avoid else frozenset()
            try:
                detoured = bfs_path(
                    state, here, self.cell, avoid=self.avoid | extra
                )
            except ValueError:
                detoured = None
            if detoured is None:
                self.avoid = frozenset()
                self.path = plain
            else:
                self.path = detoured
                # Clear only when every avoided cell is off the plain route
                # or 3+ cells ahead on it: a leader 1-2 cells ahead still
                # shares the corridor, and a same-speed follower on the
                # plain path would re-collide (the re-avoid cycle that
                # pinned two cooks in one cell). Off-route or 3+ ahead is
                # safe to follow at constant gap.
                clear = here not in self.avoid
                if clear:
                    for cell in self.avoid:
                        if cell in plain and plain.index(cell) < 3:
                            clear = False
                            break
                if clear:
                    self.avoid = frozenset()
                    self.path = plain
        else:
            self.path = bfs_path(state, here, self.cell)
        if not self.path and player.facing is not self.facing:
            # In the goal cell: the explicit face tick (the gate targets with
            # the pre-movement facing, so the face lands one tick before the
            # press) when the facing is still the wrong way; otherwise the
            # fine approach to the centre.
            return _frame(player.id, tick, face=self.facing), False
        target_cell = self.path[0] if self.path else self.cell
        target = Vec2(target_cell[0] + 0.5, target_cell[1] + 0.5)
        dist = (player.pos - target).length
        if self.path:
            if self.low is not None:
                # Low-watermark progress with a hysteresis band: a stall is
                # "worse than the best so far by the slack" (a soft-push
                # jitter swings the distance by ~0.15 per tick, so a naive
                # "did it get closer" counter resets every other tick
                # forever), and only beating the best so far by a full step
                # is progress (a waypoint hand-off jumps the distance by
                # up to ~1.0; the counter creeps to at most ~7 per
                # hand-off and resets before the next one, while a real
                # block — which never beats the watermark — accumulates
                # past the brake).
                if dist < self.low - GOTO_PROGRESS:
                    self.stuck = 0
                elif dist > self.low - GOTO_STALL_SLACK:
                    self.stuck += 1
                # else: the hysteresis band — hold the counter
            self.low = dist if self.low is None else min(self.low, dist)
        else:
            self.stuck = 0
            self.low = None
        move = (target - player.pos).normalized() if dist > 0.02 else Vec2()
        if self.path and self.stuck >= GOTO_STUCK_TICKS:
            outcome, detour, avoid = self._detour(state, here)
            if outcome == "detour":
                assert detour is not None, "the detour branch returns a path"
                self.path = detour
                self.avoid = avoid
                self.stuck = 0
                self.low = None
                target = Vec2(detour[0][0] + 0.5, detour[0][1] + 0.5)
                move = (target - player.pos).normalized()
            elif outcome == "clear":
                # No cook is on the route (or co-located): the way is
                # clear — fresh approach, no brake.
                self.stuck = 0
            else:
                move = Vec2()  # no route around: brake and let the way clear
        return _frame(player.id, tick, move=move), False

    def _detour(
        self, state: KitchenState, here: tuple[int, int]
    ) -> tuple[str, list[tuple[int, int]] | None, frozenset[tuple[int, int]]]:
        """`("detour", path, avoided)`, `("brake", None, None)` or
        `("clear", None, None)`: a route around the other cooks' cells,
        a hold (the goal is occupied, or no route around exists), or the
        all-clear (no cook is on the route — resume the approach)."""
        me = state.players[self.index]
        blocker_cells: set[tuple[int, int]] = set()
        for other in state.players:
            if other is not me:
                blocker_cells |= _body_cells(other.pos)
        blockers = frozenset(blocker_cells)
        if self.cell in blockers:
            # The goal is occupied: brake and let the blocker leave (its
            # own plan advances, so it does).
            return "brake", None, frozenset()
        plain = bfs_path(state, here, self.cell)
        co_located = here in blockers
        blocked = blockers.intersection(plain)
        if not blocked and not co_located:
            return "clear", None, frozenset()
        # A co-located blocker also blocks the first step: avoid the first
        # waypoint so the route leaves in a free direction.
        avoid = frozenset(blocked)
        if co_located:
            avoid = avoid | blockers
        if co_located and plain:
            avoid = avoid | frozenset({plain[0]})
        try:
            detour = bfs_path(state, here, self.cell, avoid=avoid)
        except ValueError:
            return "brake", None, frozenset()
        if detour == plain:
            return "brake", None, frozenset()
        return "detour", detour, blockers

    def already_there(self, player: Player) -> bool:
        """The cook stands in the band, facing the fixture — the redundant
        case the controller skips (its release tick is the face tick)."""
        here = _cell_of(player.pos)
        centre = Vec2(self.cell[0] + 0.5, self.cell[1] + 0.5)
        return (
            here == self.cell
            and (player.pos - centre).length <= GOTO_BAND
            and player.facing is self.facing
        )


    def _arrival_frame(self, player: Player, tick: int) -> tuple[PlayerInput, bool]:
        """The arrival: the explicit face tick (the gate targets with the
        pre-movement facing, so the face lands one tick before the press)."""
        return _frame(player.id, tick, face=self.facing), player.facing is self.facing


@dataclass
class _Press:
    """One press at the station (a pick, a serve, a table set-down) — one
    rising edge; the release tick follows (the controller's `release_pending`)."""

    station_id: str
    pressed: bool = False

    def tick(
        self, state: KitchenState, player: Player, tick: int
    ) -> tuple[PlayerInput, bool]:
        if not self.pressed:
            self.pressed = True
            station = _station_by_id(state, self.station_id)
            return (
                _frame(player.id, tick, action=True, face=_face_to(player, station)),
                True,
            )
        return _frame(player.id, tick, face=player.facing), True  # unreachable: done pops


@dataclass
class _PressHold:
    """The press that starts work (a set-down on a work surface, an in-place
    wash); held while the station is `WORKING` (the hold is the progress
    driver); the release tick follows when the work is done."""

    station_id: str
    pressed: bool = False

    def tick(
        self, state: KitchenState, player: Player, tick: int
    ) -> tuple[PlayerInput, bool]:
        station = _station_by_id(state, self.station_id)
        face = _face_to(player, station)
        if not self.pressed:
            self.pressed = True
            return _frame(player.id, tick, action=True, face=face), False
        if station.state is not StationState.WORKING:
            # Done — or a refusal (a raced set-down, a station that went
            # dirty in the meantime): the replan re-decides.
            return _frame(player.id, tick, face=face), True
        return _frame(player.id, tick, action=True, face=face), False


Step = _WaitUntil | _Goto | _Press | _PressHold


def _table_wash_pick(
    state: KitchenState,
    index: int,
    me_id: str,
    table_id: str,
    target_id: str,
) -> list[Step]:
    """A full hand at a dirty fixture: table the item (a marked set-down on
    the chop board, which works raw only), wash the fixture in place, pick
    the item back."""
    plan: list[Step] = []
    board_free = _board_free_predicate(table_id)
    if not board_free(state):
        plan.append(_WaitUntil(board_free, 900))
    plan.append(_Goto(state, index, table_id))
    plan.append(_Press(table_id))  # the table (no work on a non-raw set-down)
    plan.append(_Goto(state, index, target_id))
    plan.append(_PressHold(target_id))  # the in-place wash (a free hand)
    plan.append(_Goto(state, index, table_id))
    plan.append(_Press(table_id))  # the tabled item, back in the hand
    return plan


def _decide_chopper(
    state: KitchenState, index: int, me_id: str, primary_shelf: str
) -> list[Step]:
    """The chopper half of a divided two-cook run: fetch, chop, table the
    chopped on the board (a marked set-down — the board works raw only, so
    it stays a grabbable object for the cook) and repeat. The chopper never
    touches the stove; a `cook` next to it never works the board — no
    station queue, no double chop."""
    board = next(s for s in state.stations if s.kind is StationKind.CHOP_BOARD)
    primary = next(s for s in state.stations if s.id == primary_shelf)
    carried = state.players[index].carrying
    if carried is not None:
        if carried.form is ItemForm.RAW:
            plan: list[Step] = []
            board_free = _board_free_predicate(board.id)
            if not board_free(state):
                # The board is queued (a work for the other cook) or
                # loaded (a tabled object): a set-down there is a refusal,
                # so the raw goes to the shelf (a full-hand RELEASE —
                # always accepted), the board is waited on, and the raw
                # is picked back before the chop.
                plan.append(_Goto(state, index, primary.id))
                plan.append(_Press(primary.id))  # RELEASE
                plan.append(_WaitUntil(board_free, 900))
                plan.append(_Goto(state, index, primary.id))
                plan.append(_Press(primary.id))  # the raw, back in the hand
            plan.append(_Goto(state, index, board.id))
            plan.append(_PressHold(board.id))  # the set-down starts the chop
            return plan
        if carried.form is ItemForm.CHOPPED:
            # The hand-off: table the chopped on an *empty* board (a
            # set-down on a loaded surface is a refusal); the cook grabs
            # it. The board is the pass counter.
            def board_empty(s: KitchenState) -> bool:
                return (
                    _station_by_id(s, board.id).state is not StationState.WORKING
                    and not _station_by_id(s, board.id).contents
                )

            plan = []
            if not board_empty(state):
                plan.append(_WaitUntil(board_empty, 600))
            plan.append(_Goto(state, index, board.id))
            plan.append(_Press(board.id))  # the table
            return plan
    # Free hand (or an unhandled carried form, which the generic policy's
    # table-it-and-move-on covers — but the chopper's own stock first).
    if carried is None:
        if board.state is StationState.READY and board.contents:
            return [_Goto(state, index, board.id), _Press(board.id)]  # own chopped
        if primary.contents:
            return [_Goto(state, index, primary.id), _Press(primary.id)]
        return [
            _WaitUntil(
                lambda s: bool(_station_by_id(s, primary_shelf).contents), 100
            )
        ]
    return decide(state, index, primary_shelf)


def decide(
    state: KitchenState, index: int, primary_shelf: str, role: str = "any"
) -> list[Step]:
    if role == "chopper":
        return _decide_chopper(
            state, index, state.players[index].id, primary_shelf
        )
    """The policy's plan: the next full task as a list of steps (replanned
    when the list runs out — one press, one effect, then a re-look)."""
    player = state.players[index]
    me_id = player.id
    board = next(s for s in state.stations if s.kind is StationKind.CHOP_BOARD)
    stove = next(s for s in state.stations if s.kind is StationKind.STOVE)
    prep = next(s for s in state.stations if s.kind is StationKind.PREP)
    window = next(s for s in state.stations if s.kind is StationKind.SERVE)
    shelves = [s for s in state.stations if s.kind is StationKind.SPAWN_SHELF]
    primary = next(s for s in shelves if s.id == primary_shelf)
    others = [s for s in shelves if s.id != primary_shelf]

    if state.result is not None:
        return []

    carried = player.carrying
    if carried is not None:
        if carried.form is ItemForm.RAW:
            plan: list[Step] = []
            if board.state is StationState.WORKING and board.worker_id != me_id:
                plan.append(_WaitUntil(_station_free_predicate(board.id, me_id), 600))
            plan.append(_Goto(state, index, board.id))
            plan.append(_PressHold(board.id))  # the set-down starts the chop
            return plan
        if carried.form is ItemForm.CHOPPED:
            plan = []
            if _is_dirty(stove):
                plan += _table_wash_pick(state, index, me_id, board.id, stove.id)
                plan.append(_Goto(state, index, stove.id))
                plan.append(_PressHold(stove.id))  # the set-down starts the cook
                return plan
            if stove.state is StationState.WORKING and stove.worker_id != me_id:
                plan.append(_WaitUntil(_station_free_predicate(stove.id, me_id), 600))
            plan.append(_Goto(state, index, stove.id))
            plan.append(_PressHold(stove.id))
            return plan
        if carried.form is ItemForm.COOKED:
            plan = []
            if _is_dirty(prep):
                plan += _table_wash_pick(state, index, me_id, board.id, prep.id)
                plan.append(_Goto(state, index, prep.id))
                plan.append(_PressHold(prep.id))  # the press plates the dish
                return plan
            if prep.state is StationState.WORKING and prep.worker_id != me_id:
                plan.append(_WaitUntil(_station_free_predicate(prep.id, me_id), 600))
            plan.append(_Goto(state, index, prep.id))
            plan.append(_PressHold(prep.id))
            return plan
        # A `PLATED` dish: serve while the rail has an order, wait at the
        # window while it does not (the wait is free — no carry decay).
        if carried.form is ItemForm.BURNT:
            # Level 1 has no trash can: the burnt item goes to the shelf
            # (a full-hand RELEASE — always accepted, placed last, so the
            # stock picks are untouched) to free the hand.
            return [_Goto(state, index, primary.id), _Press(primary.id)]
        if state.orders:
            return [_Goto(state, index, window.id), _Press(window.id)]
        return [_WaitUntil(lambda s: bool(s.orders), 3000)]

    # A free hand. A station's contents must be usable: a ready output that
    # burnt in place (or a tabled burnt object) is never picked back.
    usable = (ItemForm.RAW, ItemForm.CHOPPED, ItemForm.COOKED, ItemForm.PLATED)
    if board.state is StationState.READY and board.contents and board.contents[0].form in usable:
        return [_Goto(state, index, board.id), _Press(board.id)]
    if stove.state is StationState.READY and stove.contents and stove.contents[0].form in usable:
        return [_Goto(state, index, stove.id), _Press(stove.id)]
    if prep.state is StationState.READY and prep.contents and prep.contents[0].form in usable:
        return [_Goto(state, index, prep.id), _Press(prep.id)]
    if board.state is StationState.IDLE and board.contents and board.contents[0].form in usable:
        return [_Goto(state, index, board.id), _Press(board.id)]  # a tabled object
    if _is_dirty(stove):
        return [_Goto(state, index, stove.id), _PressHold(stove.id)]  # the wash
    if _is_dirty(prep):
        return [_Goto(state, index, prep.id), _PressHold(prep.id)]  # the wash
    if primary.contents:
        return [_Goto(state, index, primary.id), _Press(primary.id)]
    for shelf in others:
        if shelf.contents:
            return [_Goto(state, index, shelf.id), _Press(shelf.id)]
    return [
        _WaitUntil(
            lambda s: bool(
                _station_by_id(s, primary_shelf).contents
                or any(_station_by_id(s, other.id).contents for other in others)
            ),
            100,
        )
    ]


# --- the run driver -----------------------------------------------------------


@dataclass
class CookController:
    """One scripted cook: the plan queue and the one-press-one-effect
    bookkeeping.

    `role` is the division of labour for a two-cook round: `chopper` runs
    the chopper half (fetch, chop, table the chopped on the board — the
    pass counter); `any` (the default) runs the generic full-chain policy,
    whose free-hand priority picks a tabled/ready chopped before its own
    stock — so a default cook next to a chopper runs off the hand-off.
    """

    index: int
    primary_shelf: str
    role: str = "any"
    plan: list[Step] = field(default_factory=list)
    step: Step | None = None
    release_pending: bool = True  # the first frame is a release: the gate is clean

    def next_frame(self, state: KitchenState, tick: int) -> PlayerInput:
        player = state.players[self.index]
        if self.release_pending:
            self.release_pending = False
            return _frame(player.id, tick, face=player.facing)
        if self.step is None:
            self.plan = decide(state, self.index, self.primary_shelf, self.role)
            # A Goto the cook is already in place for is redundant (the
            # release tick is the face tick): skip it so the press behind
            # it in the plan actually runs — the one-tick Goto is how the
            # controller used to loop on itself forever.
            while self.plan and isinstance(self.plan[0], _Goto):
                if not self.plan[0].already_there(player):
                    break
                self.plan.pop(0)
            if not self.plan:
                return _frame(player.id, tick, face=player.facing)
            self.step = self.plan.pop(0)
        frame, done = self.step.tick(state, player, tick)
        if done:
            self.step = None
            self.release_pending = True
        return frame


def run_round(
    state: KitchenState,
    book: RecipeBook,
    primary_shelves: Sequence[str],
    max_ticks: int,
    stop_at_verdict: bool = False,
    roles: Sequence[str] | None = None,
) -> tuple[list[KitchenState], list[list], int | None]:
    """Drive `max_ticks` substeps of the whole `step()` (or stop at the
    verdict when `stop_at_verdict`), one frame per cook per tick.

    `roles` is the per-cook division of labour (a `chopper`/`any` string
    per seat, `None` for all-generic).

    Returns:
        The per-tick states (the state after each substep), the per-tick
        event lists, and the tick the `result` first became non-`None`
        (the verdict tick — `None` when the round never ended in the window).
    """
    if len(primary_shelves) != len(state.players):
        msg = f"{len(primary_shelves)} shelves for {len(state.players)} cooks"
        raise ValueError(msg)
    if roles is not None and len(roles) != len(state.players):
        msg = f"{len(roles)} roles for {len(state.players)} cooks"
        raise ValueError(msg)
    cooks = [
        CookController(
            index, primary_shelves[index], role=(roles[index] if roles else "any")
        )
        for index in range(len(state.players))
    ]
    states: list[KitchenState] = []
    event_lists: list[list] = []
    verdict_tick: int | None = None
    current = state
    for tick in range(max_ticks):
        frames = tuple(cook.next_frame(current, tick) for cook in cooks)
        current, events = step(current, STEP_DT, frames, book)
        states.append(current)
        event_lists.append(list(events))
        if current.result is not None and verdict_tick is None:
            verdict_tick = current.tick
            if stop_at_verdict:
                break
    return states, event_lists, verdict_tick