#!/usr/bin/env python
"""One-cell Colab driver for the Golden AI 1.4 run.

In a Colab notebook on a T4, with scripts/colab/serve_share.py running on the
laptop, the whole run is two lines:

    SHARE_URL = "https://<tunnel>.trycloudflare.com"
    !wget -nv $SHARE_URL/run_v14.py -O run_v14.py && SHARE_URL=$SHARE_URL python -u run_v14.py

Installs the extras, fetches the 1.4 mix and the trainer from the laptop,
checks that real data arrived (a dead tunnel serves a 404 page, not an
error), runs the masking dry-run, then trains. The trainer uploads a
checkpoint to the laptop every 60 steps and the final adapter at the end
(build/colab_uploads/), so a reclaimed VM costs minutes. GGUF conversion is
skipped here: scripts/build_production_model.py does merge + convert +
quantise on the laptop with the exact llama.cpp build the VM serves.
"""

from __future__ import annotations

import json
import os
import pathlib
import subprocess
import sys

SHARE = os.environ.get("SHARE_URL", "").rstrip("/")
assert SHARE.startswith("https://"), "SHARE_URL must be the TUNNEL url printed by serve_share.py"

TRAIN = "sft_frc_gold_v14.jsonl"
VAL = "sft_frc_gold_v14_val.jsonl"
TRAINER = "colab_file_grounding.py"

ENV = dict(
    os.environ,
    GOLDEN_TRAIN=TRAIN,
    GOLDEN_VAL=VAL,
    GOLDEN_OUT="golden-ai-1.4-frc-gold",
    GOLDEN_SEQ=os.environ.get("GOLDEN_SEQ", "3072"),
    GOLDEN_EPOCHS=os.environ.get("GOLDEN_EPOCHS", "2"),
    GOLDEN_LR=os.environ.get("GOLDEN_LR", "1.2e-4"),
    GOLDEN_SAVE_EVERY=os.environ.get("GOLDEN_SAVE_EVERY", "60"),
    GOLDEN_UPLOAD_URL=SHARE,
    GOLDEN_SKIP_GGUF="1",
)


def sh(cmd: str, env=None) -> None:
    print("$ " + cmd, flush=True)
    rc = subprocess.run(cmd, shell=True, env=env).returncode
    if rc != 0:
        raise SystemExit(f"command failed ({rc}): {cmd}")


def main() -> int:
    sh("nvidia-smi --query-gpu=name,memory.total --format=csv")
    sh("pip -q install -U bitsandbytes peft accelerate")
    for name in (TRAIN, VAL, TRAINER):
        sh(f"wget -nv {SHARE}/{name} -O {name}")
    rows = [json.loads(l) for l in open(TRAIN, encoding="utf-8") if l.strip()]
    kinds = {}
    for r in rows:
        k = r["metadata"].get("kind") or ("file_grounding" if r["metadata"]["record_id"].startswith("file_grounding") else "anchor")
        kinds[k] = kinds.get(k, 0) + 1
    assert len(rows) > 500, f"only {len(rows)} records arrived; is serve_share.py still running?"
    trainer_text = pathlib.Path(TRAINER).read_text()
    assert "GOLDEN_TRAIN" in trainer_text and len(trainer_text) > 10000, "trainer did not arrive (a dead tunnel serves a 404 page)"
    print(f"{len(rows)} train records by kind: {kinds}", flush=True)

    sh(f"{sys.executable} {TRAINER} --dry-run", env=ENV)
    sh(f"{sys.executable} -u {TRAINER}", env=ENV)
    print("DONE: adapter uploaded to the laptop (build/colab_uploads/golden-adapter.tgz)", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())