#!/usr/bin/env python3
"""Auto-fire chain: runs once the GPU A/B job COMPLETES.

Steps (idempotent, re-runnable):
  1. Download both GPU frozen-eval JSONs + confirm checkpoints published.
  2. Regenerate the per-genre prospective figure from the GPU readouts.
  3. Refresh artifacts/paper_data.yaml A/B macros (triplet n, GPU val totals) from
     the GPU frozen evals.
  4. Launch the stats-collector job (reuse mode) to emit paper_stats_bundle.json.
  5. Commit the refreshed artifacts + figures locally.

Prints a compact completion report. Never touches the gate or frozen manifests.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
AB_REPO = "mattbitzesty/pino-pimt-representation-ab"
PY = str(ROOT / ".venv/bin/python")


def log(m): print(f"[autofire] {m}", flush=True)


def run(cmd, **kw):
    log("$ " + " ".join(str(c) for c in cmd))
    subprocess.run([str(c) for c in cmd], cwd=ROOT, check=True, **kw)


def hf_dl(name, dest):
    from huggingface_hub import hf_hub_download
    p = hf_hub_download(AB_REPO, name, repo_type="model", token=os.environ["HF_TOKEN"])
    Path(dest).write_bytes(Path(p).read_bytes())
    log(f"downloaded {name} -> {dest}")


def main() -> int:
    token = os.environ.get("HF_TOKEN")
    if not token:
        log("ERROR: HF_TOKEN not set"); return 1

    # 1. download GPU frozen evals
    evals = {}
    for arm, fname in [("morgan", "frozen_eval_morgan_gpu.json"),
                       ("openpom_256", "frozen_eval_openpom_256_gpu.json")]:
        dest = ROOT / "artifacts" / fname
        try:
            hf_dl(fname, dest)
            evals[arm] = json.loads(dest.read_text())
        except Exception as e:
            log(f"WARN: {fname} not available yet ({type(e).__name__})")
    if len(evals) < 2:
        log("ERROR: both GPU frozen evals required; aborting"); return 1

    # 2. per-genre prospective figure (GPU readouts)
    run([PY, "scripts/generate_prospective_genre_figure.py",
         "--evals", "artifacts/frozen_eval_morgan_gpu.json", "artifacts/frozen_eval_openpom_256_gpu.json",
         "--labels", "morgan", "openpom_256",
         "--output", "figures/fig_prospective_by_genre.pdf"])

    # 3. refresh paper_data A/B macros
    m, p = evals["morgan"], evals["openpom_256"]
    mt, pt = m["substitution_triplets"], p["substitution_triplets"]
    macros_patch = {
        "abTripletN": mt.get("n_triplets"),
        "abMorganTripletAcc": mt.get("accuracy"),
        "abPomTripletAcc": pt.get("accuracy"),
        "abMorganProspectiveCos": m["prospective_formulas"].get("mean_family_profile_cosine"),
        "abPomProspectiveCos": p["prospective_formulas"].get("mean_family_profile_cosine"),
        "abMorganValTotal": m.get("final_val_total"),
        "abPomValTotal": p.get("final_val_total"),
    }
    pd_path = ROOT / "artifacts/paper_data.yaml"
    txt = pd_path.read_text()
    import re
    for k, v in macros_patch.items():
        if v is None: continue
        txt = re.sub(rf"^  {k}: .*$", f"  {k}: {v}", txt, flags=re.M)
    pd_path.write_text(txt)
    log("refreshed paper_data.yaml A/B macros: " + json.dumps(macros_patch))

    # 4. commit refreshed artifacts + figures
    run(["git", "add", "artifacts/paper_data.yaml", "figures/fig_prospective_by_genre.pdf",
         "artifacts/frozen_eval_morgan_gpu.json", "artifacts/frozen_eval_openpom_256_gpu.json"])
    run(["git", "commit", "-q", "-m",
         "feat: GPU A/B readout — per-genre prospective figure + refreshed paper_data macros\n\n"
         "Auto-fired on GPU job completion. Downloaded both GPU frozen_eval JSONs (20-triplet "
         "readout), regenerated the per-genre prospective figure, and refreshed paper_data.yaml "
         "A/B macros (triplet n + accuracies + prospective cosines + val totals) to the GPU run."])

    # 5. launch stats-collector (reuse mode)
    log("launching stats-collector job (reuse mode)...")
    subprocess.run([
        "hf", "jobs", "run", "--detach", "--flavor", "t4-medium", "--timeout", "1h",
        "--secrets", "HF_TOKEN",
        "-e", "PINO_EPOCHS=20", "-e", "PINO_BATCH=32",
        "pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime", "--", "bash", "-lc",
        "set -e; pip install -q huggingface_hub; "
        "python - <<'PY2'\n"
        "import os\nfrom huggingface_hub import snapshot_download\n"
        "snapshot_download('mattbitzesty/pino-source-code', repo_type='model', local_dir='/workspace/src', token=os.environ.get('HF_TOKEN'))\n"
        "PY2\n"
        "cd /workspace/src && pip install -q -e . && python scripts/hf_stats_job.py"
    ], check=False)
    log("stats-collector launch attempted (check hf jobs ps)")

    # report
    log("=== COMPLETION REPORT ===")
    log(f"triplets: morgan {mt.get('n_correct')}/{mt.get('n_triplets')} ({mt.get('accuracy')}) | "
        f"pom {pt.get('n_correct')}/{pt.get('n_triplets')} ({pt.get('accuracy')})  [chance 0.5]")
    log(f"prospective cos: morgan {macros_patch['abMorganProspectiveCos']} | pom {macros_patch['abPomProspectiveCos']}")
    log(f"val_total: morgan {macros_patch['abMorganValTotal']} | pom {macros_patch['abPomValTotal']}")
    return 0


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