#!/usr/bin/env python3
"""One-command v4 recovery bootstrap for a FRESH Colab T4.

Context (2026-09-23): the FUKJA61G T4 died mid-real_c (uptime ~5.8h).
Everything that matters is on HF:
  * code v0.4.9 (vision-aware pooling fix)           -> repo root
  * v4 Stage B weights (loss 0.25-0.55 @ s160-164)   -> t4_assets_v4/checkpoints/hf_lora/adapter_model.pt
  * the merged 167,635-sample dataset                -> t4_assets_v3/data/dataset_final.tgz
  * HANDOVER with full RCA                           -> HANDOVER.md

This script (run INSIDE the fresh T4, after the standard code+deps bootstrap):
  1. extracts dataset_final.tgz            -> data/dataset_final  (data stage equivalent)
  2. pulls v4 hf_lora .pt                  -> checkpoints/qwenjev-s1-hf-v4/stage_a_lora + hf_lora
  3. stamps prep/data as done              -> orchestrator re-runs real_b (heads regen, s110)
     (v0.4.10: real_b is NOT skipped any more — the v4 heads file was never
      uploaded to HF and died with the FUKJA61G VM; calibrate_rlcd requires it)
  4. launches the orchestrator (GPU)       -> real_b(110) -> real_c -> ALL_DONE -> upload
  5. launches the eps0.45 collector (CPU)  -> data/traces_real3 (in parallel)

Usage (on the T4, inside /content/qwenjev after deps install):
  QWENJEV_OUT=qwenjev-s1-hf-v4 QWENJEV_ASSETS=t4_assets_v4 \
      nohup python3 -u -B scripts/v4_recover_bootstrap.py > /content/v4_recover.log 2>&1 &
"""

import json
import os
import shutil
import subprocess
import sys
import tarfile
import time
from pathlib import Path

REPO = Path("/content/qwenjev")
WORK = Path(os.environ.get("QWENJEV_WORK", "/content/qwenjev_work_v3"))
OUT = REPO / "checkpoints" / os.environ.get("QWENJEV_OUT", "qwenjev-s1-hf-v4")
ASSETS = os.environ.get("QWENJEV_ASSETS", "t4_assets_v4")
DONE = WORK / "done_v3"
HF_REPO = "tchbcb/qwenjev"

os.chdir(REPO)
token = None
for cand in (os.environ.get("HF_TOKEN"), WORK / ".hf_token"):
    if cand and str(cand).strip() and os.path.exists(str(cand)):
        token = open(str(cand)).read().strip()
        break
if not token:
    raise SystemExit("no HF_TOKEN / .hf_token found")
os.environ["HF_TOKEN"] = token

from huggingface_hub import hf_hub_download  # noqa: E402

print("[recover] 1/4 dataset_final", flush=True)
ds = REPO / "data" / "dataset_final"
if not (ds / "train" / "samples.jsonl").exists():
    tgz = hf_hub_download(HF_REPO, f"{ASSETS}/data/dataset_final.tgz"
                          if False else "t4_assets_v3/data/dataset_final.tgz")
    ds.mkdir(parents=True, exist_ok=True)
    with tarfile.open(tgz) as t:
        t.extractall(REPO)          # v0.4.10: verified archive root is
                                    # 'dataset_final' (NOT data/dataset_final)
    if not (ds / "train" / "samples.jsonl").exists():
        alt = REPO / "dataset_final"
        if (alt / "train" / "samples.jsonl").exists():
            if ds.exists():
                shutil.rmtree(ds)   # empty and/or botched nested copy —
                                    # shutil.move(dst-exists) nests src inside
            shutil.move(str(alt), str(ds))
            print("   moved dataset_final -> data/", flush=True)
print("   train rows:", sum(1 for _ in open(ds / "train" / "samples.jsonl")), flush=True)

print("[recover] 2/4 v4 stage_b weights -> stage_a_lora + hf_lora", flush=True)
pt = hf_hub_download(HF_REPO, f"{ASSETS}/checkpoints/hf_lora/adapter_model.pt")
for sub in ("stage_a_lora", "hf_lora"):
    d = OUT / sub
    d.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(pt, d / "adapter_model.safetensors")
    for f in ("adapter_config.json", "README.md"):
        try:
            src = hf_hub_download(HF_REPO, f"{ASSETS}/checkpoints/hf_lora/{f}")
            shutil.copyfile(src, d / f)
        except Exception:
            pass
meta = None
for meta_path in (f"{ASSETS}/checkpoints/hf_lora_wip/stage_b_meta.json",
                  f"{ASSETS}/checkpoints/stage_b_meta.json"):
    try:
        meta = hf_hub_download(HF_REPO, meta_path)
        shutil.copyfile(meta, OUT / "stage_b_meta.json")
        print("   stage_b_meta <-", meta_path, flush=True)
        break
    except Exception as exc:
        print("   no meta at", meta_path, ":", str(exc)[:80], flush=True)
if meta is None:
    (OUT / "stage_b_meta.json").write_text(json.dumps(
        {"wip": True,
         "note": "v4 recovery: heads missing from HF, real_b re-run "
                 "regenerates them (v0.4.10)"}, indent=2))
print("   restored ->", OUT, flush=True)

print("[recover] 3/4 stamps", flush=True)
DONE.mkdir(parents=True, exist_ok=True)
now = time.strftime("%F %T")
# v0.4.10: stamp prep+data ONLY. real_b MUST re-run: the v4 decision heads
# (OUT/qwenjev_s1.pt) were never uploaded and died with the old VM —
# calibrate_rlcd hard-fails without them ("stage B checkpoint missing").
# The s160 LoRA IS on HF, so real_b resumes from it and only the heads
# re-converge (max_steps: 110 keeps it clear of the s165 OOM wall).
for name in ("prep", "data"):
    (DONE / f"{name}.stamp").write_text(
        f"{now} (v4 recovery: prep=v2 lora re-seed skipped, "
        "data=dataset_final from HF)")
print("   stamps:", sorted(p.name for p in DONE.glob("*.stamp")), flush=True)
print("   NOTE: real_b will re-run (heads regen from s160 LoRA)", flush=True)

print("[recover] 4/4 launch orchestrator (GPU) + collector (CPU)", flush=True)
env = {**os.environ, "QWENJEV_OUT": os.environ.get("QWENJEV_OUT", "qwenjev-s1-hf-v4"),
       "QWENJEV_ASSETS": ASSETS, "PYTHONUNBUFFERED": "1"}
subprocess.Popen(["python3", "-u", "-B", "ops/t4_real25.py"],
                 cwd=str(REPO), stdout=open("/content/t4_v4.log", "w"),
                 stderr=subprocess.STDOUT, env=env,
                 start_new_session=True)
subprocess.Popen(["python3", "-u", "-B", "-m", "data.collect_selfplay",
                  "--config", "configs/collect_real_v4.yaml"],
                 cwd=str(REPO), stdout=open("/content/collect_v4.log", "w"),
                 stderr=subprocess.STDOUT, env=env,
                 start_new_session=True)
print("[recover] launched: orchestrator -> real_c; collector -> traces_real3", flush=True)
print("[recover] monitor: tail /content/t4_v4.log /content/collect_v4.log", flush=True)