import argparse
import json
import os
import shutil
import subprocess
import sys
import time
import uuid
from pathlib import Path

os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")

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)

text = (p.get("text") or "").strip()
if not text:
    print("ERROR: text is required")
    sys.exit(1)

ROOT = Path(__file__).resolve().parents[3]
MOMASK = ROOT / "vendor" / "momask"
if not (MOMASK / "mb_gen.py").exists():
    print("ERROR: momask not installed — run scripts/install_momask.sh")
    sys.exit(1)

ext = f"mb_{uuid.uuid4().hex[:10]}"  # unique run dir inside the shared vendor tree
cmd = [sys.executable, "mb_gen.py", "--gpu_id", "-1", "--ext", ext, "--text_prompt", text]
length = int(p.get("length", -1))
if length > 0:
    cmd += ["--motion_length", str(length)]

t0 = time.time()
r = subprocess.run(cmd, cwd=str(MOMASK), capture_output=True, text=True)
gen_dir = MOMASK / "generation" / ext
files = sorted(gen_dir.rglob("*")) if gen_dir.exists() else []
bvh = [f for f in files if f.suffix == ".bvh"]
if r.returncode != 0 or not bvh:
    print(r.stdout[-1000:])
    print(r.stderr[-1500:], file=sys.stderr)
    print("ERROR: momask generation failed")
    sys.exit(1)

outdir = Path(a.outdir)
slug = "".join(c if c.isalnum() else "_" for c in text[:40]).strip("_")
n = 0
for f in files:
    if f.suffix == ".bvh":
        # _ik = foot-contact cleaned; that's the one to retarget
        kind = "ik" if "_ik" in f.stem else "raw"
        shutil.copy(f, outdir / f"motion_{slug}_{kind}.bvh")
        n += 1
    elif f.suffix == ".mp4" and "_ik" in f.stem:
        shutil.copy(f, outdir / f"motion_{slug}_preview.mp4")
        n += 1
shutil.rmtree(gen_dir, ignore_errors=True)
print(f"done in {time.time() - t0:.0f}s: {n} files for '{text[:60]}'", flush=True)