#!/bin/bash
# Pull raw records, logs and chain messages from the cluster nodes and merge the E1_core Llama shard into runs/raw/E1_core.jsonl
# (idempotent: shard records are appended only if their attempt_id is not already present).
set -e
cd "$(dirname "$0")/.."
mkdir -p runs/raw/remote runs/logs/remote runs/raw/chains
rsync -aq M3U96a:~/wp30-worker/wp30/runs/raw/E1_core.jsonl runs/raw/remote/E1_core_llama_m3u96a.jsonl
rsync -aq M3U96a:~/wp30-worker/wp30/runs/logs/ runs/logs/remote/m3u96a/
for f in ED_hosted_core E4_core E5_core E3_core E3_core_sender; do rsync -aq M3U96b:~/wp30-worker/wp30/runs/raw/$f.jsonl runs/raw/ 2>/dev/null || true; done
rsync -aq M3U96b:~/wp30-worker/wp30/runs/raw/chains/ runs/raw/chains/ 2>/dev/null || true
rsync -aq M3U96b:~/wp30-worker/wp30/runs/logs/ runs/logs/remote/m3u96b/
python3 - <<'PY'
import json
seen = set()
for l in open("runs/raw/E1_core.jsonl"):
    seen.add(json.loads(l)["attempt_id"])
n = 0
with open("runs/raw/E1_core.jsonl", "a") as out:
    for l in open("runs/raw/remote/E1_core_llama_m3u96a.jsonl"):
        r = json.loads(l)
        if r["attempt_id"] not in seen:
            r["hardware_id"] = "m3u96a-m3ultra-96g"; out.write(json.dumps(r, ensure_ascii=False) + "\n"); seen.add(r["attempt_id"]); n += 1
print("merged", n, "new Llama-shard records; total", len(seen))
# relabel hardware for records produced on M3U96b by code that still carried the laptop constant
for f in ["E4_core", "E5_core", "E3_core", "E3_core_sender"]:
    p = f"runs/raw/{f}.jsonl"
    try:
        rows = [json.loads(l) for l in open(p)]
    except FileNotFoundError:
        continue
    ch = 0
    for r in rows:
        if r.get("hardware_id") == "laptop-m5max-48g":
            r["hardware_id"] = "m3u96b-m3ultra-96g"; ch += 1
    with open(p, "w") as out:
        for r in rows: out.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f, "relabelled", ch)
PY