"""Scale up: run many randomized episodes per task and keep only the ones that succeeded.

Each episode is a fresh seed, so the scene randomization differs every time; the scripted solver
plays it and the task's own conditions decide whether it counts. Failed episodes are discarded
(or kept under --keep-failures for debugging), which is what makes this a demonstration
generator rather than a batch runner.

    python scripts/generate.py --episodes 20                      # every task
    python scripts/generate.py --tasks grape_box,stack_blocks --episodes 50
    python scripts/generate.py --tasks passing --episodes 100     # only tasks known to work
    python scripts/generate.py --episodes 20 --resume             # skip episodes already done

Writes:
    <out>/<task>/ep<seed>.mp4        video, successes only unless --keep-failures
    <out>/<task>/ep<seed>.json       per-episode result + the exact randomized placements
    <out>/manifest.json              running totals and success rate per task

One episode is one simulator launch, so this is slow and long-running by nature; it prints a
line per episode and is safe to interrupt and --resume.
"""
import argparse
import json
import os
import re
import subprocess
import sys
import time

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

ap = argparse.ArgumentParser()
ap.add_argument("--tasks", default="all",
                help="'all', 'passing', or a comma-separated list of task names")
ap.add_argument("--episodes", type=int, default=10, help="episodes per task")
ap.add_argument("--start-seed", type=int, default=100,
                help="seeds run from here upward; keep clear of the 0-9 used for development")
ap.add_argument("--out", default="data")
ap.add_argument("--keep-failures", action="store_true")
ap.add_argument("--resume", action="store_true", help="skip episodes whose json already exists")
ap.add_argument("--timeout", type=int, default=1500)
ap.add_argument("--python", default=sys.executable)
a = ap.parse_args()


def registered():
    out = subprocess.run([a.python, "scripts/yam_task.py", "--list"], cwd=REPO,
                         capture_output=True, text=True).stdout
    return [l.split()[0] for l in out.splitlines()[1:] if l.strip() and not l.startswith(("name", "-"))
            and len(l.split()) > 2]


def known_passing():
    """Tasks whose most recent development run succeeded -- the ones worth spending seeds on."""
    good = []
    for t in registered():
        log = os.path.join(REPO, "outputs", "tasks", f"v3_{t}.log")
        if os.path.exists(log) and "EPISODE_RESULT: SUCCESS" in open(log, errors="ignore").read():
            good.append(t)
    return good


tasks = (registered() if a.tasks == "all" else
         known_passing() if a.tasks == "passing" else
         [t.strip() for t in a.tasks.split(",") if t.strip()])
if not tasks:
    raise SystemExit("no tasks selected")

out_root = os.path.join(REPO, a.out)
os.makedirs(out_root, exist_ok=True)
manifest_path = os.path.join(out_root, "manifest.json")
manifest = json.load(open(manifest_path)) if os.path.exists(manifest_path) else {}

env = dict(os.environ)
env.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")
env.setdefault("PYTHONUNBUFFERED", "1")

print(f"[gen] {len(tasks)} task(s) x {a.episodes} episodes, seeds "
      f"{a.start_seed}..{a.start_seed+a.episodes-1} -> {out_root}", flush=True)
t_start = time.time()

for task in tasks:
    tdir = os.path.join(out_root, task)
    os.makedirs(tdir, exist_ok=True)
    rec = manifest.setdefault(task, {"attempted": 0, "succeeded": 0, "episodes": []})

    for i in range(a.episodes):
        seed = a.start_seed+i
        meta_p = os.path.join(tdir, f"ep{seed}.json")
        if a.resume and os.path.exists(meta_p):
            continue
        vid = os.path.join(tdir, f"ep{seed}.mp4")
        t0 = time.time()
        try:
            p = subprocess.run(
                [a.python, "scripts/yam_task.py", "--task", task, "--seed", str(seed),
                 "--video", vid,
                 "--kit_args=--/rtx/verifyDriverVersion/enabled=false"],
                cwd=REPO, env=env, capture_output=True, text=True, timeout=a.timeout)
            log = p.stdout+p.stderr
        except subprocess.TimeoutExpired:
            log = "TIMEOUT"
        ok = "EPISODE_RESULT: SUCCESS" in log
        # the randomized placements the env printed -- this is the episode's actual scene
        m = re.search(r"placements=(\{.*?\})\s*$", log, re.M)
        meta = {"task": task, "seed": seed, "success": ok,
                "seconds": round(time.time()-t0, 1),
                "placements": m.group(1) if m else None,
                "conditions": re.findall(r"^\[env\]   (PASS|FAIL)  (.+)$", log, re.M)}
        json.dump(meta, open(meta_p, "w"), indent=1)
        if not ok and not a.keep_failures and os.path.exists(vid):
            os.remove(vid)
        rec["attempted"] += 1
        rec["succeeded"] += int(ok)
        rec["episodes"].append({"seed": seed, "success": ok})
        json.dump(manifest, open(manifest_path, "w"), indent=1)
        print(f"[gen] {task:<18} seed {seed:<4} {'ok  ' if ok else 'FAIL'} "
              f"{meta['seconds']:>5.1f}s   ({rec['succeeded']}/{rec['attempted']})", flush=True)

print(f"\n[gen] done in {(time.time()-t_start)/60:.1f} min")
tot_a = sum(r["attempted"] for r in manifest.values())
tot_s = sum(r["succeeded"] for r in manifest.values())
for t, r in sorted(manifest.items()):
    if r["attempted"]:
        print(f"  {t:<20} {r['succeeded']:>4}/{r['attempted']:<4} "
              f"{100*r['succeeded']/r['attempted']:5.1f}%")
print(f"  {'TOTAL':<20} {tot_s:>4}/{tot_a:<4} "
      f"{(100*tot_s/tot_a if tot_a else 0):5.1f}%   manifest -> {manifest_path}")