#!/usr/bin/env python3
"""Bootstrap the NEW T4 instance end-to-end (replaces the timed-out one).

Usage: T4_URL=<new-tunnel> python3 scripts/new_t4_bootstrap.py

Steps (all idempotent):
  1. sanity-check the envelope
  2. upload qwenjev_code.tgz + unpack into /content/qwenjev
  3. pip install arc-agi, upgrade torchao, write HF token
  4. launch remote re-collection (24 games x 10 eps, nohup)
  5. launch ops/t4_real25.py orchestrator (nohup) — its prep stage runs
     immediately (HF pull v2 adapter), data stage rebuilds dataset_remote
     if the collector is still running... NOTE: we wait for collection by
     launching the orchestrator AFTER collection finishes; use mode=lite to
     train on the local-only dataset right away instead.
"""
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ops"))
import t4exec  # noqa: E402

TOKEN = os.environ.get("HF_TOKEN")  # never hardcode; pass via env
CODE_TGZ = os.environ.get("CODE_TGZ", "/tmp/qwenjev_code.tgz")
MODE = os.environ.get("MODE", "full")  # full = wait for re-collection


def run(code, timeout=120):
    r = t4exec.execute(code, timeout=timeout)
    out = r.get("stdout", "")
    if not out.strip():
        raise RuntimeError(f"empty response: { {k: str(v)[:120] for k, v in r.items()} }")
    return out


def main():
    assert TOKEN, "export HF_TOKEN=<hf write token> first"
    print("== 1. envelope sanity ==")
    print(run("print('envelope ok')").strip())

    print("== 2. code upload ==")
    t4exec.put(CODE_TGZ, "/content/qwenjev_code.tgz", timeout=300)
    run("""
import tarfile, os, shutil
os.makedirs('/content/qwenjev', exist_ok=True)
with tarfile.open('/content/qwenjev_code.tgz') as t:
    t.extractall('/content/qwenjev')
os.remove('/content/qwenjev_code.tgz')
print('unpacked:', os.path.isdir('/content/qwenjev/ops'))
""")

    print("== 3. deps ==")
    run("""
import subprocess
r = subprocess.run(['bash','-c',
  'pip install -q arc-agi 2>&1 | tail -1; pip install -q -U torchao 2>&1 | tail -1; '
  'pip install -q -U bitsandbytes 2>&1 | tail -1; pip install -q peft 2>&1 | tail -1'],
  capture_output=True, text=True, timeout=300)
print('deps done')
""", timeout=420)
    run(f"""
import os
os.makedirs('/content/qwenjev_work_v3/logs', exist_ok=True)
os.makedirs('/content/qwenjev_work_v3/done_v3', exist_ok=True)
open('/content/qwenjev_work_v3/.hf_token','w').write({TOKEN!r})
print('dirs+token ok')
""")

    print("== 4. remote re-collection ==")
    run("""
import subprocess
p = subprocess.Popen(['bash','-c',
  'cd /content/qwenjev && nohup python3 -u data/collect_selfplay.py '
  '--config configs/collect_real24.yaml > /content/collect_v3.log 2>&1 < /dev/null &'],
  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
print('collect launched')
""")

    print("== 5. orchestrator handoff ==")
    if MODE == "full":
        # chain daemon: polls the collect log, launches the orchestrator
        # the moment collection finishes — fully unattended.
        run("""
import subprocess
p = subprocess.Popen(['bash','-c',
  'cd /content/qwenjev && nohup python3 -u scripts/remote_chain.py > /content/chain.log 2>&1 < /dev/null &'],
  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
print('chain daemon on')
""")
    elif MODE == "lite":
        # train on the local-only dataset (HF tgz) right away; the data
        # stage of t4_real25.py reuses dataset_remote if present else
        # rebuilds it — for lite mode we let it build from traces (still
        # collecting) is NOT ok; so we pre-seed dataset_remote from HF.
        run("""
import subprocess
p = subprocess.Popen(['bash','-c',
  'cd /content/qwenjev && python3 - <<EOF > /tmp/seed.log 2>&1\\n'
  'import tarfile\\n'
  'from huggingface_hub import hf_hub_download\\n'
  'tgz = hf_hub_download("tchbcb/qwenjev", "t4_assets_v3/data/dataset_v3_local.tgz")\\n'
  'with tarfile.open(tgz) as t: t.extractall("data/")\\n'
  'import shutil, os\\n'
  'shutil.move("data/dataset_v3", "data/dataset_remote")\\n'
  'print("seeded")\\nEOF'
  ' && cp /tmp/seed.log /content/seed_done.log'],
  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
print('seeding local dataset as dataset_remote (async)')
""")
    else:
        run("""
import subprocess
p = subprocess.Popen(['bash','-c',
  'cd /content/qwenjev && nohup python3 -u -B ops/t4_real25.py > /content/t4_v3.log 2>&1 < /dev/null &'],
  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
print('orchestrator launched (lite)')
""")
    # watchdog: auto control-eval after ALL_DONE (both modes)
    run("""
import subprocess
p = subprocess.Popen(['bash','-c',
  'cd /content/qwenjev && nohup python3 -u scripts/watchdog_eval_v3.py > /content/watchdog.log 2>&1 < /dev/null &'],
  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
print('watchdog on')
""")
    print("BOOTSTRAP_DONE")


if __name__ == "__main__":
    main()