undercooked/tests/sim/test_result_payload.py

"""`sim/round_state.result_payload` — the body of the `round` message (PLAN §7).

The task's contract, pinned one rule at a time:

* **JSON-safety.** The payload is the wire body: plain JSON types only
  (`None`, `bool`, `int`, finite `float`, `str`, list, dict) — no `Vec2`,
  no enum, no dataclass, no NaN/inf. A recursive `assert_json_safe` helper
  walks the whole payload and names the first leak by path;
* **The key set.** The top-level keys, exactly; the `counters` projection,
  exactly `served` / `missed` / `wasted` (a never-written counter reads as
  0, an unrelated counter — `trashed` — does not leak in); the per-player
  mapping, exactly `id` / `name` / `score` / `dishes_served` /
  `missed_contrib`;
* **The values.** `result` is the enum's string (or `None` before the
  verdict), `stars` is `sim/round_state.stars` applied to the same state,
  `missed_contrib` is the documented display-only even split of the round's
  `missed` count, `players` in seat order;
* **Stability.** Two identical states — and the same state re-read after a
  whole `SETTLE_S` settle window — project identically, byte for byte
  (`json.dumps` with sorted keys): the payload is a function of the state
  at the verdict alone, no matter when the `round` message is sent.
"""

from __future__ import annotations

import copy
import json
import math
from collections.abc import Mapping
from dataclasses import is_dataclass, replace
from enum import Enum
from typing import cast

import pytest

from undercooked.sim.grid import Grid
from undercooked.sim.level import Level, LevelDef
from undercooked.sim.model import (
    Facing,
    Hazard,
    HazardKind,
    Item,
    ItemForm,
    ItemKind,
    KitchenState,
    Order,
    Player,
    RoundResult,
    Station,
    StationKind,
    StationState,
    TileKind,
    Vec2,
)
from undercooked.sim.recipe import RecipeBook
from undercooked.sim.round_state import SETTLE_S, result_payload, stars
from undercooked.sim.state_build import SeatSpec, build_state
from undercooked.sim.step import STEP_DT, step

pytestmark = pytest.mark.unit

# --- the JSON-safety helper -----------------------------------------------------------


def assert_json_safe(value: object, path: str = "root") -> None:
    """Recursively assert `value` is plain JSON; name the first leak by path.

    The leak classes the guard exists for: a `Vec2` (or any dataclass) that
    serialised as an object of the wrong shape, an enum that slipped through
    the string check (a `StrEnum` *is* a `str`), a NaN/inf float that would
    be invalid JSON on the wire, and a non-`str` mapping key.
    """
    if value is None:
        return
    if isinstance(value, Enum):
        msg = f"{path}: enum {type(value).__name__} leaked into the payload"
        raise AssertionError(msg)
    if is_dataclass(value) and not isinstance(value, type):
        msg = f"{path}: dataclass {type(value).__name__} leaked into the payload"
        raise AssertionError(msg)
    if isinstance(value, bool | str):
        return
    if isinstance(value, int):
        return
    if isinstance(value, float):
        if not math.isfinite(value):
            msg = f"{path}: non-finite float {value!r} is invalid JSON"
            raise AssertionError(msg)
        return
    if isinstance(value, (list, tuple)):
        for index, item in enumerate(value):
            assert_json_safe(item, f"{path}[{index}]")
        return
    if isinstance(value, Mapping):
        for key, item in value.items():
            if not isinstance(key, str):
                msg = f"{path}: mapping key {key!r} is not a str"
                raise AssertionError(msg)
            assert_json_safe(item, f"{path}.{key}")
        return
    msg = f"{path}: {type(value).__name__} is not a JSON type"
    raise AssertionError(msg)


# --- narrowing helpers (the payload is `Mapping[str, object]`) ---------------------------


def players_of(payload: Mapping[str, object]) -> list[dict[str, object]]:
    """The `players` list, narrowed for the type checker (and asserted a list)."""
    raw = payload["players"]
    assert isinstance(raw, list)
    return cast("list[dict[str, object]]", raw)


def counters_of(payload: Mapping[str, object]) -> Mapping[str, int]:
    """The `counters` mapping, narrowed (and asserted a mapping of ints)."""
    raw = payload["counters"]
    assert isinstance(raw, Mapping)
    return cast("Mapping[str, int]", raw)


# --- fixtures ---------------------------------------------------------------------------


def make_state(
    *,
    result: RoundResult | None = None,
    counters: dict[str, int] | None = None,
) -> KitchenState:
    """A fully-populated kitchen: mixed tiles, a `WORKING` station, two
    cooks (seat order deliberately reversed), a live order, a hazard, set-down
    marks — and counters with an extra key (`trashed`) to prove the
    projection."""
    grid = Grid(
        width=8,
        height=8,
        tiles=(
            TileKind.WALL,
            TileKind.FLOOR,
            TileKind.COUNTER,
            TileKind.VOID,
        )
        * 16,
    )
    a = Player(
        id="b-token",
        pos=Vec2(1.5, 1.5),
        facing=Facing.DOWN,
        seat=1,
        display_name="Bee",
        score=25,
        dishes_served=1,
        carrying=Item(kind=ItemKind.TOMATO, form=ItemForm.PLATED, quality=0.9),
    )
    b = Player(
        id="a-token",
        pos=Vec2(2.5, 3.5),
        facing=Facing.UP,
        seat=0,
        display_name="Ash",
        score=15,
        dishes_served=2,
    )
    station = Station(
        id="stove_0",
        kind=StationKind.STOVE,
        x=4.0,
        y=4.0,
        state=StationState.WORKING,
        progress=0.42,
        duration=6.0,
        contents=(Item(kind=ItemKind.TOMATO, form=ItemForm.COOKING),),
        worker_id="b-token",
    )
    return KitchenState(
        level_id="05_payload",
        seed=4242,
        grid=grid,
        tick=321,
        stations=(station,),
        # seat order deliberately reversed: the payload must sort it.
        players=(a, b),
        orders=(
            Order(id=7, recipe_id="soup", time_left=12.5, deadline=420.0, value=40),
        ),
        order_pool=("soup",),
        order_concurrent=1,
        order_spawn_min=80,
        order_spawn_max=120,
        order_deadline=240,
        order_penalty=10,
        next_order_at=500,
        next_order_id=8,
        hazards=(
            Hazard(
                id="fire_0",
                kind=HazardKind.FIRE,
                pos=Vec2(4.0, 4.0),
                life=3.5,
                station_id="stove_0",
            ),
        ),
        score=40,
        target_score=60,
        time_limit=100.0,
        time_left=60.0,
        streak=2,
        result=result,
        counters=(
            {"served": 3, "missed": 5, "wasted": 2, "longest_streak": 4, "trashed": 1}
            if counters is None
            else counters
        ),
        item_marks={"stove_0": (("a-token", 300),)},
        rng_state=987654,
        accumulator=0.025,
    )


def book() -> RecipeBook:
    return RecipeBook.model_validate(
        {
            "recipes": {
                "soup": {
                    "id": "soup",
                    "name": "Tomato Soup",
                    "base_score": 40,
                    "ingredients": [{"kind": "tomato", "form": "raw"}],
                    "steps": [{"kind": "serve", "station": "serve"}],
                }
            }
        }
    )


# --- JSON-safety -----------------------------------------------------------------------


def test_the_payload_is_json_safe_with_zero_leaks() -> None:
    payload = result_payload(make_state(result=RoundResult.WON))
    assert_json_safe(payload)  # the recursive guard, over the whole body


def test_the_payload_survives_a_json_round_trip() -> None:
    payload = result_payload(make_state(result=RoundResult.WON))
    decoded = json.loads(json.dumps(payload))
    assert decoded == payload  # floats round-trip exactly (repr-based)


def test_the_payload_never_leaks_a_vec2_an_enum_or_a_dataclass() -> None:
    # The state is thick with the leak classes: Vec2 positions, StrEnum tile
    # kinds, a Hazard dataclass, an Item in a hand. None of them may reach
    # the payload — walk it and name the first one that does.
    state = make_state(result=RoundResult.ABORTED)
    payload = result_payload(state)
    assert_json_safe(payload)

    def types(value: object) -> set[str]:
        found: set[str] = set()
        if isinstance(value, Mapping):
            for item in value.values():
                found |= types(item)
        elif isinstance(value, (list, tuple)):
            for item in value:
                found |= types(item)
        else:
            found.add(type(value).__name__)
        return found

    leaked = types(payload) - {
        "NoneType",
        "bool",
        "int",
        "float",
        "str",
        "list",
        "dict",
    }
    assert leaked == set()


# --- the key set and the values ---------------------------------------------------------


def test_the_payload_has_the_exact_top_level_key_set() -> None:
    payload = result_payload(make_state())
    assert set(payload) == {
        "result",
        "score",
        "target",
        "stars",
        "time_left",
        "counters",
        "longest_streak",
        "players",
        "level_id",
        "seed",
        "ticks",
    }


def test_the_scalar_types_are_plain_json_types() -> None:
    payload = result_payload(make_state(result=RoundResult.WON))
    assert isinstance(payload["result"], str)
    assert isinstance(payload["score"], int)
    assert isinstance(payload["target"], int)
    assert isinstance(payload["stars"], int)
    assert isinstance(payload["time_left"], float)
    assert isinstance(payload["longest_streak"], int)
    assert isinstance(payload["level_id"], str)
    assert isinstance(payload["seed"], int)
    assert isinstance(payload["ticks"], int)
    assert isinstance(payload["players"], list)


def test_the_counters_projection_is_exactly_served_missed_wasted() -> None:
    payload = result_payload(make_state())
    counters = counters_of(payload)
    assert set(counters) == {"served", "missed", "wasted"}
    assert counters["served"] == 3
    assert counters["missed"] == 5
    assert counters["wasted"] == 2
    # The state's `trashed` counter does not leak into the projection.
    assert "trashed" not in counters


def test_a_never_written_counter_reads_as_zero() -> None:
    payload = result_payload(make_state(counters={}))
    assert payload["counters"] == {"served": 0, "missed": 0, "wasted": 0}
    assert payload["longest_streak"] == 0


def test_result_and_stars_follow_the_round_state() -> None:
    assert result_payload(make_state())["result"] is None  # not ended yet
    assert result_payload(make_state())["stars"] == 0
    for outcome, label in (
        (RoundResult.WON, "won"),
        (RoundResult.LOST, "lost"),
        (RoundResult.ABORTED, "aborted"),
    ):
        payload = result_payload(make_state(result=outcome))
        assert payload["result"] == label
        assert payload["stars"] == stars(make_state(result=outcome))
        assert payload["stars"] == (3 if outcome is RoundResult.WON else 0)


def test_the_per_player_mapping_has_the_exact_key_set_and_types() -> None:
    payload = result_payload(make_state())
    players = players_of(payload)
    assert len(players) == 2
    for entry in players:
        assert set(entry) == {
            "id",
            "name",
            "score",
            "dishes_served",
            "missed_contrib",
        }
        assert isinstance(entry["id"], str)
        assert isinstance(entry["name"], str)
        assert isinstance(entry["score"], int)
        assert isinstance(entry["dishes_served"], int)
        assert isinstance(entry["missed_contrib"], int)


def test_the_players_are_in_seat_order() -> None:
    # The state's players are deliberately reversed (seat 1 first); the
    # payload sorts by seat, then id.
    payload = result_payload(make_state())
    players = players_of(payload)
    assert [entry["id"] for entry in players] == ["a-token", "b-token"]
    assert [entry["name"] for entry in players] == ["Ash", "Bee"]


@pytest.mark.parametrize(
    ("missed", "seats", "each"),
    [
        (3, 2, 1),  # floor(3 / 2)
        (5, 2, 2),  # floor(5 / 2)
        (0, 4, 0),  # nothing to spread
        (7, 1, 7),  # one cook takes the whole estimate
    ],
)
def test_missed_contrib_is_the_display_only_even_split(
    missed: int, seats: int, each: int
) -> None:
    state = make_state(counters={"missed": missed})
    state = replace(
        state,
        players=tuple(
            Player(id=f"p{index}", seat=index, display_name=f"Cook {index}")
            for index in range(seats)
        ),
    )
    payload = result_payload(state)
    players = players_of(payload)
    assert len(players) == seats
    for entry in players:
        assert entry["missed_contrib"] == each


def test_a_round_with_no_cooks_projects_an_empty_player_list() -> None:
    state = make_state()
    state = replace(state, players=())
    payload = result_payload(state)
    assert payload["players"] == []


# --- stability --------------------------------------------------------------------------


def test_the_payload_is_stable_across_two_identical_states() -> None:
    first = result_payload(make_state(result=RoundResult.WON))
    second = result_payload(make_state(result=RoundResult.WON))
    assert first == second
    assert json.dumps(first, sort_keys=True) == json.dumps(second, sort_keys=True)


def test_the_payload_is_stable_across_a_settle_window() -> None:
    # The body the room sends at the buzzer and the one it could send
    # `SETTLE_S` later are the same bytes in everything the verdict fixes —
    # score, stars, time_left, counters, the players. Only `ticks` moves
    # through the window (it is the settle's purpose), so it is compared
    # apart.
    current = replace(make_state(), time_left=60.0, score=60, target_score=60)
    current, _events = step(current, STEP_DT)  # the early win lands
    at_verdict = result_payload(current)
    for _ in range(round(SETTLE_S / STEP_DT)):
        current, _events = step(current, STEP_DT)
    at_settle_end = result_payload(current)
    ticks_verdict = at_verdict["ticks"]
    ticks_settle = at_settle_end["ticks"]
    assert isinstance(ticks_verdict, int)
    assert isinstance(ticks_settle, int)
    assert ticks_settle == ticks_verdict + round(SETTLE_S / STEP_DT)
    for key in at_verdict:
        if key != "ticks":
            assert at_verdict[key] == at_settle_end[key], key
    without_ticks = {key: value for key, value in at_verdict.items() if key != "ticks"}
    without_ticks_end = {
        key: value for key, value in at_settle_end.items() if key != "ticks"
    }
    assert json.dumps(without_ticks, sort_keys=True) == json.dumps(
        without_ticks_end, sort_keys=True
    )


# --- build_state wiring ------------------------------------------------------------------


TILES: dict[str, dict[str, str]] = {
    ".": {"kind": "floor"},
    "#": {"kind": "wall"},
    "S": {"kind": "floor", "station": "spawn_shelf", "item": "tomato"},
    "V": {"kind": "floor", "station": "serve"},
}

GRID = [
    "#######",
    "#S...V#",
    "#.....#",
    "#######",
]


def make_level() -> Level:
    data: dict[str, object] = {
        "id": "05_payload_build",
        "name": "Payload Kitchen",
        "order": 1,
        "time_limit": 45.0,
        "target_score": 30,
        "difficulty": 1,
        "grid": copy.deepcopy(GRID),
        "tiles": copy.deepcopy(TILES),
        "spawns": [[1, 1]],
        "orders": {
            "pool": ["soup"],
            "concurrent": 1,
            "spawn_interval": [4, 6],
            "deadline": 30,
        },
    }
    return Level.build(LevelDef.model_validate(data))


def test_a_built_round_projects_the_seat_name_and_the_limit() -> None:
    level = make_level()
    seats = (
        SeatSpec(player_id="p0", seat=0, display_name="Reese"),
        SeatSpec(player_id="p1", seat=1, display_name="Moe"),
    )
    first = build_state(level, seats, 7, book())
    second = build_state(level, seats, 7, book())
    for state in (first, second):
        payload = result_payload(state)
        assert [entry["name"] for entry in players_of(payload)] == ["Reese", "Moe"]
        assert payload["time_left"] == 45.0
        assert payload["result"] is None
    # Two identical builds, stepped a little, project byte-identically.
    stepped_first = first
    stepped_second = second
    for _ in range(10):
        stepped_first, _events = step(stepped_first, STEP_DT, book=book())
        stepped_second, _events = step(stepped_second, STEP_DT, book=book())
    assert json.dumps(result_payload(stepped_first), sort_keys=True) == json.dumps(
        result_payload(stepped_second), sort_keys=True
    )