"""HF Jobs SFT launcher — exact mirror of working Colab cells.

No TRL OpenEnv, no env server, no GRPO. Pure SFT on 611 traces.
Same code that works on Colab T4. H200 just runs it ~10× faster.
"""
from __future__ import annotations

import os
import subprocess
import sys

WORK = "/tmp/opsguard"
HUB_REPO = "sai1906/opsguard-sft"


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


def main():
    if not os.environ.get("HF_TOKEN"):
        raise SystemExit("HF_TOKEN required")

    sh([sys.executable, "-m", "pip", "install", "-q", "-U",
        "transformers>=4.46", "trl>=0.18", "peft>=0.13", "bitsandbytes>=0.44",
        "datasets", "accelerate>=1.0", "huggingface_hub", "matplotlib", "networkx>=3"])

    sh(f"rm -rf {WORK}")
    sh(["git", "clone", "https://huggingface.co/spaces/sai1906/opsguard", WORK])
    os.chdir(WORK)
    sys.path.insert(0, WORK)

    from huggingface_hub import login, HfApi
    login(token=os.environ["HF_TOKEN"])
    HfApi(token=os.environ["HF_TOKEN"]).create_repo(HUB_REPO, repo_type="model", exist_ok=True)
    print(f"[preflight] repo {HUB_REPO} ready", flush=True)

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

    MODEL = "Qwen/Qwen2.5-7B-Instruct"
    print(f"[model] loading {MODEL} 4bit...", flush=True)
    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                              bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
    tok = AutoTokenizer.from_pretrained(MODEL)
    if tok.pad_token_id is None:
        tok.pad_token_id = tok.eos_token_id
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, quantization_config=bnb, torch_dtype=torch.bfloat16, device_map="auto",
    )
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    lc = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, bias="none",
                    target_modules=["q_proj","k_proj","v_proj","o_proj",
                                    "gate_proj","up_proj","down_proj"],
                    task_type="CAUSAL_LM")
    model = get_peft_model(model, lc)
    model.print_trainable_parameters()

    import json
    from datasets import Dataset
    rows = [json.loads(l) for l in open("data/sft_traces.jsonl")]
    print(f"[data] {len(rows)} SFT traces loaded", flush=True)
    texts = [r["prompt"] + "\n\nACTION:\n" + r["completion"] + tok.eos_token for r in rows]
    ds = Dataset.from_list([{"text": t} for t in texts])

    from trl import SFTConfig, SFTTrainer
    import inspect

    def safe_kwargs(cls, kw):
        sig = inspect.signature(cls).parameters
        return {k: v for k, v in kw.items() if k in sig}

    raw = dict(
        output_dir="/tmp/opsguard-sft",
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        num_train_epochs=2,
        learning_rate=1e-4,
        warmup_ratio=0.03,
        logging_steps=5,
        save_strategy="epoch",
        save_total_limit=1,
        bf16=True,
        max_seq_length=2048,
        dataset_text_field="text",
        report_to="none",
        push_to_hub=False,
    )
    cfg = SFTConfig(**safe_kwargs(SFTConfig, raw))
    trainer = SFTTrainer(model=model, train_dataset=ds, args=cfg, processing_class=tok)
    trainer.train()

    print("[push] saving + pushing to Hub...", flush=True)
    model.save_pretrained("/tmp/opsguard-sft-lora")
    tok.save_pretrained("/tmp/opsguard-sft-lora")
    HfApi(token=os.environ["HF_TOKEN"]).upload_folder(
        folder_path="/tmp/opsguard-sft-lora", repo_id=HUB_REPO, repo_type="model",
    )
    print(f"DONE: pushed to https://huggingface.co/{HUB_REPO}", flush=True)


if __name__ == "__main__":
    main()