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

A toy durable-execution runtime: activity(fn, *args) returns the recorded
result for step k from a JSON log if there is one, and otherwise runs fn and
records it. A workflow of three activities is killed (os._exit) after the
second; the second start replays 1 and 2 from the log and runs only 3. Then
a workflow that branches on time.time() is replayed and diverges.
"""
import argparse, json, os, subprocess, sys, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

RUNTIME = r'''
import json, os, sys, time
LOG = sys.argv[1]; MODE = sys.argv[2]; CRASH = sys.argv[3] == "crash"
log = json.load(open(LOG)) if os.path.exists(LOG) else []
step = [0]
def activity(fn, *a):
    k = step[0]; step[0] += 1
    if k < len(log):
        if log[k]["name"] != fn.__name__:
            print(f"DIVERGED at step {k}: log has {log[k]['name']}, workflow called {fn.__name__}"); sys.exit(2)
        print(f"replayed {fn.__name__} from the log"); return log[k]["result"]
    r = fn(*a); log.append({"name": fn.__name__, "result": r})
    json.dump(log, open(LOG, "w")); return r
def charge(x): print("RAN charge"); return x * 2
def ship(x): print("RAN ship"); return x + 1
def email(x): print("RAN email"); return x
def refund(x): print("RAN refund"); return -x
a = activity(charge, 21)
if MODE == "clock" and int(time.time() * 10) % 2: b = activity(refund, a)       # branches on the clock: not deterministic
else: b = activity(ship, a)
if CRASH: os._exit(1)                                                          # killed after the second activity
c = activity(email, b); print("done", c)
'''
d = tempfile.mkdtemp(); src = os.path.join(d, "wf.py"); open(src, "w").write(RUNTIME)
def start(log, mode, crash):
    r = subprocess.run([sys.executable, "-u", src, log, mode, "crash" if crash else "ok"], capture_output=True, text=True)
    return r.returncode, [l for l in r.stdout.splitlines()]

rows = []
log = os.path.join(d, "det.json")
for i, crash in enumerate((True, False)):
    code, out = start(log, "det", crash)
    rows.append({"workflow": "deterministic", "start": i + 1, "exit": code, "ran": ", ".join(l[4:] for l in out if l.startswith("RAN")) or "-",
                 "replayed": ", ".join(l.split()[1] for l in out if l.startswith("replayed")) or "-", "outcome": "killed" if code == 1 else out[-1]})
diverged = 0
for attempt in range(10):                                              # try until the clock picks the other branch on replay
    log = os.path.join(d, f"clock{attempt}.json")
    code1, out1 = start(log, "clock", True); time.sleep(0.05); code2, out2 = start(log, "clock", False)
    if code2 == 2:
        diverged += 1
        rows.append({"workflow": "branches on time.time()", "start": 2, "exit": code2, "ran": ", ".join(l[4:] for l in out2 if l.startswith("RAN")) or "-",
                     "replayed": ", ".join(l.split()[1] for l in out2 if l.startswith("replayed")) or "-", "outcome": out2[-1]})
        break
summary = {"clock_attempts": attempt + 1, "diverged": "yes" if diverged else "no"}
for r in rows: print(f"{r['workflow']:<24} start {r['start']}: ran [{r['ran']}], replayed [{r['replayed']}] -> {r['outcome']}")

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