#!/usr/bin/env python3
"""
Little Fig — GPU Benchmark Runner (HF Jobs compatible)
Clones repo, installs deps, runs benchmark, saves results to HF Hub.
"""

import os, sys, subprocess, json, time

# Clone and install
print("[SETUP] Cloning Little Fig...", flush=True)
if not os.path.exists("/app/littlefig"):
    subprocess.check_call(["git", "clone", "https://github.com/ticketguy/littlefig.git", "/app/littlefig"])

print("[SETUP] Installing deps...", flush=True)
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q",
    "transformers", "accelerate", "peft", "bitsandbytes",
    "datasets", "sentencepiece", "protobuf", "psutil", "numpy", "torch"])
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", "/app/littlefig[train]"])

# Run the benchmark
print("[SETUP] Running benchmark...", flush=True)
os.chdir("/app/littlefig")
result = subprocess.run([sys.executable, "benchmark/benchmark_gpu.py"],
    capture_output=False, text=True)

# Upload results if they exist
if os.path.exists("/app/benchmark_results.json"):
    from huggingface_hub import HfApi
    api = HfApi()
    api.upload_file(
        path_or_fileobj="/app/benchmark_results.json",
        path_in_repo=f"benchmark_results_{int(time.time())}.json",
        repo_id="ticketguy/littlefig-benchmarks",
    )
    print("[DONE] Results uploaded to HF Hub", flush=True)
else:
    print("[WARN] No results file found", flush=True)