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

ROOT = Path(__file__).resolve().parents[3]
BRUSH = ROOT / "bin" / "brush"

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 colmap dataset")
    sys.exit(1)
dataset = Path(a.input[0])
if not BRUSH.exists():
    print(f"ERROR: brush not installed at {BRUSH}. Run scripts/install_brush.sh")
    sys.exit(1)
# brush wants the dataset root containing images/ + sparse/
if not (dataset / "images").exists():
    inner = dataset / "colmap_dataset"
    if (inner / "images").exists():
        dataset = inner

outdir = Path(a.outdir)
export = outdir / "export"
export.mkdir(parents=True, exist_ok=True)
steps = int(p.get("total_steps", 30000))

cmd = [
    str(BRUSH), str(dataset.resolve()),
    "--total-steps", str(steps),
    "--max-splats", str(p.get("max_splats", 10000000)),
    "--sh-degree", str(p.get("sh_degree", 3)),
    "--max-resolution", str(p.get("max_resolution", 1920)),
    "--seed", str(p.get("seed", 42)),
    "--export-every", str(steps),          # export once at the end
    "--export-path", str(export.resolve()),
    "--export-name", "splat.ply",
]
print("+", " ".join(cmd), flush=True)
res = subprocess.run(cmd, stdout=sys.stdout, stderr=subprocess.STDOUT)
if res.returncode != 0:
    sys.exit(res.returncode)

plys = sorted(export.rglob("*.ply"))
if not plys:
    print("ERROR: brush produced no .ply splat")
    sys.exit(1)
(outdir / "result.json").write_text(json.dumps({"outputs": [
    {"path": str(plys[-1]), "name": f"{dataset.name}_splat.ply",
     "meta": {"kind": "splat", "tool": "brush", "steps": steps}}]}))
print("done:", plys[-1], flush=True)