"""Phase 1: write evidence bundle files into ev/out and print name/sha/size."""
import json, glob, os, hashlib, shutil
O = "ev/out"
def w(name, obj):
    p = f"{O}/{name}"
    open(p, "w").write(obj if isinstance(obj, str) else json.dumps(obj, indent=1, default=str))
for f in glob.glob("ev/pred_*.jsonl") + ["ev/scored-observations.jsonl"]:
    shutil.copy(f, f"{O}/{os.path.basename(f)}")
bundle = {}
for d in sorted(glob.glob("ev/rank*")) + sorted(glob.glob("ev/pilot_*")):
    if not os.path.isdir(d): continue
    a = os.path.basename(d); bundle[a] = {}
    for f in sorted(glob.glob(f"{d}/*.json")):
        if "logprobs" in f: continue
        bundle[a][os.path.basename(f)] = json.load(open(f))
    bundle[a]["group_rewards_per_group_call"] = [json.loads(l) for l in open(f"{d}/group_rewards.jsonl")] if os.path.exists(f"{d}/group_rewards.jsonl") else None
w("init-provenance-bundle.json", bundle)
lp = {os.path.relpath(f, "ev"): json.load(open(f)) for f in sorted(glob.glob("ev/*/initial_logprobs*.json")) + sorted(glob.glob("ev/probe_initial_logprobs_*.json"))}
lp["fixed_batch"] = json.load(open("ev/fixed_batch.json"))
w("fixed-batch-and-initial-logprobs.json", lp)
w("probe-results.json", json.load(open("ev/probe_results.json")))
tm = {}
for d in sorted(glob.glob("runs/*")):
    a = os.path.basename(d)
    tm[a] = {"metrics": [{k: v for k, v in json.loads(l).items() if not k.startswith("time/") or k == "time/total"} for l in open(f"{d}/metrics.jsonl")],
             "checkpoints": [json.loads(l) for l in open(f"{d}/checkpoints.jsonl")] if os.path.exists(f"{d}/checkpoints.jsonl") else [],
             "recipe_config": json.load(open(f"{d}/config.json")) if os.path.exists(f"{d}/config.json") else None}
w("training-metrics.json", tm)
scores = {os.path.basename(f): {k: v for k, v in json.load(open(f)).items() if k != "provenance"} for f in sorted(glob.glob("ev/score_*.json"))}
scores["selection.json"] = json.load(open("ev/selection.json"))
scores["primary"] = json.load(open("ev/report_parts.json"))["primary"]
scores["per_arm"] = json.load(open("ev/report_parts.json"))["arms"]
w("eval-scores-and-selection.json", scores)
for s in ["common.py", "probe.py", "train_arm.py", "eval_gsm8k.py", "select_lr.py", "build_report.py", "auto_sel.sh", "assemble.py"]:
    shutil.copy(s, f"{O}/{s}")
w("code-revision.json", {"tinker": "0.37.0", "tinker-cookbook": "git+https://github.com/thinking-machines-lab/tinker-cookbook.git@f79fdb2ba4cd967ef2807ba43097c52985f9f17f",
    "research_exchange": "modal-projects/research-exchange@77745cd (validators only)", "datasets": "5.1.0", "transformers": "5.22.1",
    "base_model": "Qwen/Qwen2.5-3B@3aab1f1954e9cc14eb9509a215f9e5ca08227a9b", "gsm8k": "openai/gsm8k main @740312add88f781978c0658806c59bc2815b9866",
    "spindle": "shared Spindle at $TINKER_BASE_URL (no fork)", "frozen_code_sha256": open("ev/frozen_code_sha256.txt").read(),
    "train_arm.py_note": "train_arm.py was edited after the first launch to (a) skip save_weights_external and (b) write *_restore1 files on resume; final sha in this bundle"})
logs = {}
for f in ["ev/probe_attempt1.log", "ev/probe.log"]:
    logs[f] = "".join(l for l in open(f) if "FutureWarning" not in l and "return func" not in l and "httpx" not in l)[-20000:]
for f in sorted(glob.glob("logs/rank32_*_attempt1.log")):
    t = open(f).read(); i = t.find("Traceback"); logs[f] = t[i:i + 6000] if i >= 0 else t[-3000:]
w("logs-excerpts.json", logs)
for f in sorted(os.listdir(O)):
    b = open(f"{O}/{f}", "rb").read(); print(f, len(b), hashlib.sha256(b).hexdigest())