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

ROOT = Path(__file__).resolve().parents[3]

# fork checkout resolution: vendored path first, then the per-box conventional
# clones (m3 staging, m1 home) so existing installs work without re-vendoring.
CANDIDATES = [
    ROOT / "vendor" / "trellis2-mlx",
    Path.home() / "Documents" / "trellis2-mlx-staging" / "trellis2-jourloy",
    Path.home() / "trellis2-mlx",
]
FORK = next((c for c in CANDIDATES if (c / "scripts" / "generate_asset.py").exists()), None)

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 FORK is None:
    print("ERROR: trellis-2-mrp-mlx not found. Run scripts/install_trellis2_mlx.sh")
    sys.exit(1)

py = FORK / ".venv" / "bin" / "python"
if not py.exists():
    print(f"ERROR: no venv at {FORK}/.venv — run its scripts/setup_macos.sh")
    sys.exit(1)

outdir = Path(a.outdir)
cmd = [
    str(py), str(FORK / "scripts" / "generate_asset.py"),
    str(Path(a.input[0]).resolve()),
    "--output-dir", str(outdir.resolve()),
    "--backend", "mlx-experimental",
    "--baker", str(p.get("baker", "vertex")),
    "--pipeline-type", str(p.get("pipeline_type", "1024_cascade")),
    "--seed", str(p.get("seed", 0)),
    "--texture-size", "2048",
    "--force",
]
if p.get("steps"):
    cmd += ["--steps", str(p["steps"])]

env = os.environ.copy()
env["HF_HUB_DISABLE_XET"] = "1"

print("+", " ".join(cmd), flush=True)
res = subprocess.run(cmd, cwd=str(FORK), stdout=sys.stdout,
                     stderr=subprocess.STDOUT, env=env)
if res.returncode != 0:
    sys.exit(res.returncode)

glbs = sorted(outdir.rglob("candidate_pbr.glb")) or sorted(outdir.rglob("*.glb"))
if not glbs:
    print("ERROR: no GLB produced")
    sys.exit(1)
stem = Path(a.input[0]).stem
outputs = [{"path": str(glbs[0]), "name": f"trellis2mlx_{stem}.glb",
            "meta": {"tool": "trellis2_mlx", "baker": p.get("baker", "vertex")}}]
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
print("done:", glbs[0], flush=True)