import os, sys, time, subprocess, json, re
from huggingface_hub import HfApi

NAMESPACE = "GAInTech"
REPO_ID = "GAInTech/feather-pretrain-checkpoints"
IMAGE = "GAInTech/feather-a10g-large-runtime"
TPS_FLOOR = 150000
BEST_BPB_VAL = 0.8726 # prod9 A10G champion (bpb, not ppl); Cluster E baseline was 2.9696
RUN_LABEL = "long-horizon-stabilized"

def get_active_job():
    try:
        r = subprocess.run(["hf", "jobs", "ps", "--namespace", NAMESPACE], capture_output=True, text=True)
        lines = r.stdout.strip().splitlines()
        for ln in lines:
            if "RUNNING" in ln or "PENDING" in ln:
                return ln.split()[0]
    except: pass
    return None

def monitor_job(job_id):
    try:
        r = subprocess.run(["hf", "jobs", "logs", "--namespace", NAMESPACE, job_id, "--tail", "100"], capture_output=True, text=True)
        out = r.stdout
        # Extract last step TPS and BPB
        metrics = re.findall(r"step=(\d+).*bpb=([\d\.]+).*tps=(\d+)", out)
        if not metrics: return True # Wait more
        
        last_step, last_bpb, last_tps = metrics[-1]
        last_step, last_bpb, last_tps = int(last_step), float(last_bpb), int(last_tps)
        
        print(f"[Guardian] Job {job_id} | Step {last_step} | BPB {last_bpb} | TPS {last_tps}")
        
        # Audit 2026-05-13: Kill if NaNs detected in log
        if "nan" in out.lower():
            print(f"[Guardian] NaNs detected in log. Killing.")
            return False

        # Audit 2026-05-13: allow 20 steps of data warmup before TPS floor
        if last_tps < TPS_FLOOR and last_step > 20:
            print(f"[Guardian] TPS {last_tps} below floor {TPS_FLOOR}. Killing.")
            return False
            
        # Refined trajectory check: kill if step 50 is still worse than champion
        if last_bpb > (BEST_BPB_VAL * 1.2) and last_step > 50:
            print(f"[Guardian] BPB {last_bpb} significantly worse than champion {BEST_BPB_VAL}. Killing.")
            return False
            
        return True
    except: return True

def launch_resume(source_job_id):
    print(f"[Guardian] Launching resume from {source_job_id}...")
    env = os.environ.copy()
    env["FEATHER_HF_OWNER"] = "GAInTech"
    env["FEATHER_HF_JOB_NAMESPACE"] = "GAInTech"
    env["FEATHER_HF_SPACE_REPO"] = IMAGE
    env["FEATHER_HF_USE_SPACE_IMAGE"] = "1"
    env["FEATHER_HF_SKIP_UPLOAD"] = "1"
    env["HYDRA_RESUME_JOB_ID"] = source_job_id
    env["HYDRA_RESUME_CKPT_NAME"] = "pretrain_final.pt"
    # Match the champion's engram and retina arch exactly
    env["HYDRA_ENGRAM_N_COLUMNS"] = "1024"
    env["HYDRA_CONTRASTIVE_RANK"] = "0"
    # Full optimizer restore enabled
    env["HYDRA_RESUME_RESET_OPTIMIZER"] = "0"
    env["HYDRA_MATRIX_LR"] = "0.04"
    env["HYDRA_USE_NEMOTRON"] = "1"
    env["HYDRA_LOCAL_SHARDS_ONLY"] = "0"
    
    cmd = [sys.executable, "scripts/launch_feather_hf_job.py"]
    subprocess.run(cmd, env=env)

def main():
    job_id = get_active_job()
    if not job_id:
        # Resume from the actual champion
        launch_resume("6a01d522317220dbbd1a7a6a")
    else:
        is_healthy = monitor_job(job_id)
        if not is_healthy:
            subprocess.run(["hf", "jobs", "cancel", "--namespace", NAMESPACE, job_id])
            # Next tick will relaunch

if __name__ == "__main__":
    main()