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

ROOT = Path(__file__).resolve().parents[3]
VENDOR = ROOT / "vendor" / "trellis-mac"

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 / "generate.py").exists():
    print(f"ERROR: trellis-mac not installed at {VENDOR}. Run scripts/install_trellis_mac.sh")
    sys.exit(1)

outdir = Path(a.outdir)
out_stem = outdir / f"trellis2_{Path(a.input[0]).stem}"

cmd = [
    sys.executable, "generate.py", str(Path(a.input[0]).resolve()),
    "--pipeline-type", str(p.get("pipeline_type", "1024")),
    "--texture-size", str(p.get("texture_size", 2048)),
    "--seed", str(p.get("seed", 0)),
    "--output", str(out_stem.resolve()),
]
if p.get("no_texture"):
    cmd.append("--no-texture")

env = os.environ.copy()
# xet chunked downloader fails on these BFL/Meta repos ("Unable to parse string
# as hex hash value") — force the reliable HTTP path (same fix as flux_local).
env["HF_HUB_DISABLE_XET"] = "1"
# guard the torch + compiled-kernel OpenMP collision like sf3d does.
env["KMP_DUPLICATE_LIB_OK"] = "TRUE"
env.setdefault("OMP_NUM_THREADS", "1")
env.setdefault("MKL_NUM_THREADS", "1")

print("+", " ".join(cmd), flush=True)
print("(first run downloads TRELLIS.2-4B + dinov3 + RMBG-2.0 weights)", flush=True)
res = subprocess.run(cmd, cwd=str(VENDOR), stdout=sys.stdout, stderr=subprocess.STDOUT, env=env)
if res.returncode != 0:
    sys.exit(res.returncode)

# generate.py writes <output>.glb (and often a .obj); collect any glb under outdir
glbs = sorted(outdir.rglob("*.glb"))
if not glbs:
    print("ERROR: trellis-mac produced no .glb")
    sys.exit(1)
outputs = [{"path": str(glbs[0]), "name": f"trellis2_{Path(a.input[0]).stem}.glb",
            "meta": {"tool": "trellis_mac"}}]
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
print("done:", glbs[0], flush=True)