"""HF Jobs launcher — boots env locally + runs training in same container.

Optimized 3-round coevolution-light:
  Round 0: SFT warmstart (Unsloth, 1 epoch on 611 traces)
  Round 1: GRPO defender (train_grpo.py, 100 steps) vs scripted env adversary
  Round 2: GRPO defender (train_grpo.py, 100 steps) vs scripted env adversary (E2-E4)
  Round 3: Final eval + push LoRA to Hub

Total: ~3 × 30min on a100-large = 90min, ~$3.75.

Usage on HF Jobs:
    hf jobs uv run --flavor a100-large \
        --with "trl>=0.18,unsloth,peft,bitsandbytes,openenv-core,vllm,matplotlib,networkx,datasets" \
        --secrets HF_TOKEN \
        https://huggingface.co/spaces/sai1906/opsguard/raw/main/scripts/hf_launcher.py
"""
from __future__ import annotations

import os
import subprocess
import sys
import time
from pathlib import Path


REPO_ZIP = "https://huggingface.co/spaces/sai1906/opsguard/resolve/main"
WORK = Path("/tmp/opsguard")


def sh(cmd, check=True, **kw):
    print(f"[sh] {cmd if isinstance(cmd, str) else ' '.join(cmd)}", flush=True)
    return subprocess.run(cmd, shell=isinstance(cmd, str), check=check, **kw)


def clone_repo():
    if WORK.exists():
        sh(["rm", "-rf", str(WORK)])
    WORK.mkdir(parents=True, exist_ok=True)
    sh(["git", "clone", "https://huggingface.co/spaces/sai1906/opsguard", str(WORK)])
    os.chdir(WORK)
    sys.path.insert(0, str(WORK))


def install_deps():
    sh([sys.executable, "-m", "pip", "install", "-q",
        "trl>=0.18", "peft>=0.13", "bitsandbytes>=0.44", "openenv-core[core]>=0.2.2",
        "matplotlib", "networkx>=3", "datasets", "huggingface_hub", "accelerate>=1.0",
        "transformers>=4.46", "jmespath", "fastapi", "uvicorn", "requests"])


def boot_env_server() -> subprocess.Popen:
    print("[env] starting opsguard env server on 0.0.0.0:8001...", flush=True)
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8001"],
        cwd=str(WORK),
        env={**os.environ, "PYTHONPATH": str(WORK)},
    )
    for _ in range(60):
        time.sleep(2)
        try:
            import urllib.request
            urllib.request.urlopen("http://0.0.0.0:8001/health", timeout=2).read()
            print("[env] healthy", flush=True)
            return proc
        except Exception:
            continue
    raise RuntimeError("env server did not become healthy within 120s")


def sft_warmstart():
    print("\n=== SFT WARM-START ===", flush=True)
    sh([sys.executable, "scripts/sft_warmstart.py",
        "--model", "Qwen/Qwen2.5-3B-Instruct",
        "--episodes", "32",
        "--output-dir", "/tmp/opsguard-sft",
        "--epochs", "1",
        "--train-only",
        ])


def grpo_round(round_idx: int, num_steps: int, sft_adapter: str | None, hub_repo: str):
    print(f"\n=== ROUND {round_idx} — GRPO DEFENDER ({num_steps} steps) ===", flush=True)
    args = [
        sys.executable, "scripts/train_grpo.py",
        "--model", "unsloth/Qwen2.5-7B-Instruct-bnb-4bit",
        "--env-url", "http://0.0.0.0:8001",
        "--num-steps", str(num_steps),
        "--num-generations", "4",
        "--max-steps-per-episode", "30",
        "--scenarios", "E2_social_eng_buildup", "E3_compromised_maintainer", "E4_multi_vector",
        "--no-vllm",
        "--max-completion-length", "512",
        "--output-dir", f"/tmp/opsguard-grpo-r{round_idx}",
    ]
    if sft_adapter:
        args.extend(["--sft-adapter", sft_adapter])
    if hub_repo:
        args.extend(["--hub-repo", hub_repo])
    sh(args)


def final_eval():
    print("\n=== FINAL EVAL: trained LoRA vs baseline policies ===", flush=True)
    sh([sys.executable, "scripts/run_baseline_eval.py",
        "--policies", "keyword_security_triager", "memory_aware", "debate", "ensemble_voting",
        "--scenarios", "E2_social_eng_buildup", "E3_compromised_maintainer",
                       "E4_multi_vector", "E5_jia_tan_saga",
        "--seeds", "0", "1", "2",
        "--out", "/tmp/eval_post_train"])


def push_outputs():
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ.get("HF_TOKEN"))
    try:
        api.create_repo("sai1906/opsguard-grpo", repo_type="model", exist_ok=True)
    except Exception as e:
        print(f"[push] create_repo: {e}", flush=True)
    for src, dest in [
        ("/tmp/opsguard-grpo-r2", "checkpoints/r2"),
        ("/tmp/eval_post_train", "eval_post_train"),
    ]:
        if Path(src).exists():
            try:
                api.upload_folder(folder_path=src, repo_id="sai1906/opsguard-grpo",
                                   path_in_repo=dest, repo_type="model")
                print(f"[push] uploaded {src} -> {dest}", flush=True)
            except Exception as e:
                print(f"[push] upload failed for {src}: {e}", flush=True)


def main():
    if not os.environ.get("HF_TOKEN"):
        raise SystemExit("HF_TOKEN required as secret")
    clone_repo()
    install_deps()
    try:
        from huggingface_hub import HfApi as _HfApi
        _HfApi(token=os.environ["HF_TOKEN"]).create_repo(
            "sai1906/opsguard-sft", repo_type="model", exist_ok=True
        )
        print("[preflight] hub repo ready", flush=True)
    except Exception as _e:
        print(f"[preflight] create_repo: {_e}", flush=True)
    sft_train_and_push()
    final_eval()
    push_outputs()


def sft_train_and_push():
    print("\n=== SFT TRAINING (Qwen 7B 4bit + LoRA) ===", flush=True)
    sh([sys.executable, "scripts/sft_warmstart.py",
        "--model", "unsloth/Qwen2.5-7B-Instruct-bnb-4bit",
        "--episodes", "32",
        "--output-dir", "/tmp/opsguard-sft",
        "--epochs", "2",
        "--train-only",
        "--push-to-hub", "sai1906/opsguard-sft",
        ])


if __name__ == "__main__":
    main()