# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Replay determinism.

A tiny exactly-once pipeline: a loop over 10,000 events updates a dictionary
and checkpoints (state + position) every 500 events. It is crashed at a
random position and restarted from the checkpoint, 50 times, first with a
pure computation and then with one that reads random.random(). Each final
state is compared with a crash-free run. Then the impure output is written as
an upsert keyed by event id, and replayed events are counted.
"""
import argparse, copy, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--events", type=int, default=10_000, help="Events.")
p.add_argument("--every", type=int, default=500, help="Checkpoint interval.")
p.add_argument("--trials", type=int, default=50, help="Crash-and-restart trials.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = random.Random(args.seed)
events = [{"id": i, "user": f"u{rng.randrange(50)}", "amount": rng.randrange(1, 100)} for i in range(args.events)]
def step(state, e, impure):
    bonus = random.random() if impure else 0.0                             # the hidden non-determinism
    state[e["user"]] = state.get(e["user"], 0) + e["amount"] + bonus
def run(impure, crash_at=None, sink=None):
    ckpt = ({}, 0); state, pos = {}, 0; replayed = 0
    while pos < len(events):
        e = events[pos]; step(state, e, impure)
        if sink is not None:
            if e["id"] in sink: replayed += 1
            sink[e["id"]] = round(state[e["user"]], 6)                     # upsert keyed by event id
        pos += 1
        if pos % args.every == 0: ckpt = (copy.deepcopy(state), pos)
        if crash_at is not None and pos == crash_at:
            state, pos = copy.deepcopy(ckpt[0]), ckpt[1]; crash_at = None   # crash: restart from the checkpoint
    return state, replayed
rows = []
for impure in (False, True):
    random.seed(1); base, _ = run(impure)
    differ = 0
    for t in range(args.trials):
        crash = rng.randrange(1, args.events); random.seed(1)               # same seed as the crash-free run: only the replay differs
        st, _ = run(impure, crash)
        if any(abs(st.get(k, 0) - v) > 1e-9 for k, v in base.items()): differ += 1
    rows.append({"computation": "random.random() inside" if impure else "pure", "trials": args.trials, "final_state_differs": differ})
sink = {}; random.seed(7); _, replayed = run(True, crash_at=args.events // 2 + 123, sink=sink)
summary = {"sink_rows": len(sink), "events": args.events, "replayed_writes": replayed}
for r in rows: print(f"{r['computation']:<24} final state differs from the crash-free run in {r['final_state_differs']} of {r['trials']} crashes (no error raised)")
print(f"upsert keyed by event id: {replayed} replayed writes, {len(sink):,} rows for {args.events:,} events")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"computations": rows}}, f)