#!/usr/bin/env python3
"""Reusable high-quality splat generator: 524K gaussians (with fallbacks) + 32 steps.
Usage: python3 gen_hq.py <slug> [<slug> ...]  (reads images/<slug>.jpg -> splats/<slug>.ply)"""
import sys, pathlib, time
import hf_call

ROOT = "https://vast-ai-triposplat.hf.space"
IMG  = pathlib.Path(__file__).parent / "images"
OUT  = pathlib.Path(__file__).parent / "splats"
OUT.mkdir(exist_ok=True)

GAUSSIANS = [524288, 393216, 327680, 262144]
STEPS = 32

for slug in sys.argv[1:]:
    src = IMG / f"{slug}.jpg"
    dest = OUT / f"{slug}.ply"
    if not src.exists():
        print(f"[miss] {slug}: no image", flush=True); continue
    ok = False
    for ng in GAUSSIANS:
        for attempt in range(1, 4):
            try:
                print(f"[up  ] {slug} (ng={ng}, try {attempt}) ...", flush=True)
                fd = hf_call.upload(ROOT, src)
                print(f"[gen ] {slug}: {ng} gaussians, {STEPS} steps ...", flush=True)
                data = hf_call.call(ROOT, "generate", [fd, 42, STEPS, 3.0, ng, "ply"], timeout=2400)
                hf_call.download(data[1], dest, space_root=ROOT)
                print(f"[ok  ] {slug} -> {dest} ({dest.stat().st_size//1024} kB, {ng} gaussians)", flush=True)
                ok = True; break
            except Exception as e:
                print(f"[err ] {slug} ng={ng} attempt {attempt}: {e}", flush=True)
                time.sleep(12)
        if ok: break
    if not ok:
        print(f"[FAIL] {slug}", flush=True)

print("DONE")