import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
VENDOR = ROOT / "vendor" / "stable-fast-3d"

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not a.input:
    print("ERROR: no input image")
    sys.exit(1)
if not (VENDOR / "run.py").exists():
    print(f"ERROR: SF3D not installed at {VENDOR}. Run scripts/install_sf3d.sh")
    sys.exit(1)

outdir = Path(a.outdir)
work = outdir / "sf3d_out"
work.mkdir(parents=True, exist_ok=True)

cmd = [
    sys.executable, "run.py", str(Path(a.input[0]).resolve()),
    "--device", p.get("device", "mps"),
    "--output-dir", str(work.resolve()),
    "--texture-resolution", str(p.get("texture_resolution", 1024)),
    "--foreground-ratio", str(p.get("foreground_ratio", 0.85)),
    "--remesh_option", p.get("remesh_option", "none"),
]
env = os.environ.copy()
env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
print("+", " ".join(cmd), flush=True)
res = subprocess.run(cmd, cwd=str(VENDOR), env=env, stdout=sys.stdout, stderr=subprocess.STDOUT)
if res.returncode != 0:
    sys.exit(res.returncode)

glbs = sorted(work.rglob("*.glb"))
if not glbs:
    print("ERROR: SF3D produced no .glb")
    sys.exit(1)
name = f"sf3d_{Path(a.input[0]).stem}.glb"
(outdir / "result.json").write_text(json.dumps(
    {"outputs": [{"path": str(glbs[0]), "name": name, "meta": {"tool": "sf3d"}}]}))
print("done:", glbs[0], flush=True)