"""One-screen status of the pipeline and a training run.
Usage: source env.sh && $TA_PY scripts/status.py [run_dir]   (default: $TA_DATA/runs/main)"""
import glob
import json
import os
import sys

from tiny_agent.text import DATA

run = sys.argv[1] if len(sys.argv) > 1 else f"{DATA}/runs/main"
pl = f"{DATA}/logs/pipeline.log"
if os.path.exists(pl):
    print("pipeline:")
    for l in open(pl).read().splitlines()[-12:]:
        print("  " + l)
lj = f"{run}/log.jsonl"
if not os.path.exists(lj):
    sys.exit(f"no log yet at {lj}")
recs = [json.loads(l) for l in open(lj)]
last = recs[-1]
args = json.load(open(f"{run}/meta.json")) if os.path.exists(f"{run}/meta.json") else {}
print(f"\nrun {run}")
print(f"  step {last['step']}  tokens {last['tokens']/1e9:.3f}B  loss {last['loss']}  lr_mult {last['lr_mult']}  "
      f"tok/s {last['tok_s']}  elapsed {last['elapsed_min']/60:.1f}h")
vals = [r for r in recs if "val" in r]
for r in vals[-4:]:
    print(f"  val @ {r['tokens']/1e9:.2f}B: " + " ".join(f"{k}={v:.3f}" for k, v in sorted(r["val"].items())))
snaps = sorted(glob.glob(f"{run}/snap_*.pt"))
print(f"  snapshots: {', '.join(os.path.basename(s) for s in snaps[-6:]) or 'none'}")