scripts/generate_demo.py

"""
Pre-record a Polis run into data/demo_run.json.

Run with a real key for rich behaviour:
    OPENAI_API_KEY=sk-... python -m scripts.generate_demo --ticks 40

Run with no key to produce a deterministic mock demo (still fully playable):
    python -m scripts.generate_demo --ticks 40

The frontend plays this file so the Hugging Face Space always has something to
show, even at zero API budget.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from backend.world import World, LOCATIONS
from backend.llm import llm

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "data" / "demo_run.json"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ticks", type=int, default=40)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    world = World.bootstrap(seed=args.seed)
    snapshots = []

    # scripted world events to make the story move
    injects = {8: "a stranger arrives at the harbor carrying a locked chest",
               18: "a sudden storm floods the market stalls",
               30: "news spreads of a festival planned for the plaza"}

    for t in range(1, args.ticks + 1):
        if t in injects:
            world.inject_event(injects[t])
        snap = world.step()
        snapshots.append(snap)
        print(f"tick {t:>2} | events {len(snap['events']):>2} | "
              f"spent ${llm.ledger.spent_usd:.4f} | live={llm.live}")

    payload = {
        "meta": {
            "ticks": args.ticks, "seed": args.seed, "live": llm.live,
            "budget": llm.ledger.as_dict(),
            "locations": LOCATIONS,
            "personas": [{"id": a.id, "name": a.name, "age": a.age, "role": a.role,
                          "traits": a.traits, "color": a.color, "home": a.home}
                         for a in world.agents],
            "injects": injects,
        },
        "snapshots": snapshots,
        "final_agents": [a.snapshot() for a in world.agents],
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(payload, indent=None))
    print(f"\nwrote {OUT}  ({OUT.stat().st_size // 1024} KB)")


if __name__ == "__main__":
    main()