#!/usr/bin/env python3
"""Turn each monument image into a 3D gaussian splat (.ply) via the
VAST-AI/TripoSplat Gradio space."""
import pathlib, sys, 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)

SLUGS = ["tour-eiffel", "louvre-pyramid", "arc-de-triomphe",
         "sacre-coeur", "notre-dame", "obelisque"]

# allow running a subset:  python3 gen_splats.py tour-eiffel obelisque
if len(sys.argv) > 1:
    SLUGS = sys.argv[1:]

for slug in SLUGS:
    src = IMG / f"{slug}.jpg"
    dest = OUT / f"{slug}.ply"
    if dest.exists():
        print(f"[skip] {slug}", flush=True)
        continue
    if not src.exists():
        print(f"[miss] {slug} image not found", flush=True)
        continue
    ok = False
    for attempt in range(1, 5):
        try:
            print(f"[up  ] {slug} uploading image (try {attempt}) ...", flush=True)
            filedata = hf_call.upload(ROOT, src)
            print(f"[gen ] {slug} generating splat (this can take a few min) ...", flush=True)
            data = hf_call.call(ROOT, "generate", [
                filedata,   # image
                42,         # seed
                20,         # steps
                3.0,        # guidance_scale
                262144,     # num_gaussians
                "ply",      # output_format
            ], timeout=1800)
            # returns (preprocessed_image, ply_file, download_file, info_string)
            ply = data[1]
            hf_call.download(ply, dest, space_root=ROOT)
            print(f"[ok  ] {slug} -> {dest} ({dest.stat().st_size//1024} kB)", flush=True)
            ok = True
            break
        except Exception as e:
            print(f"[err ] {slug} attempt {attempt}: {e}", flush=True)
            time.sleep(15)
    if not ok:
        print(f"[FAIL] {slug} gave up after retries", flush=True)

print("DONE")