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]
ARDY = Path.home() / "Documents" / "ardy"
ARDY2BVH = Path.home() / "Documents" / "ardy2bvh" / "ardy2bvh.py"
if not (ARDY / "scripts" / "generate.py").exists():
    print("ERROR: ardy repo not found at ~/Documents/ardy")
    sys.exit(1)

env = dict(os.environ)
env["LOCAL_CACHE"] = "true"
env["TEXT_ENCODERS_DIR"] = str(ROOT / "data" / "text_encoders")
# The adapters in TEXT_ENCODERS_DIR pin the ungated Llama-3 mirror; never let a
# stale ~/.cache/huggingface/token break those public downloads.
env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
env["TOKENIZERS_PARALLELISM"] = "false"

stem = f"mb_{uuid.uuid4().hex[:10]}"
workdir = Path(a.outdir) / "_work"
workdir.mkdir(parents=True, exist_ok=True)

cmd = [
    sys.executable, "scripts/generate.py", text,
    "--model", p.get("model", "core"),
    "--duration", str(float(p.get("duration", 5.0))),
    "--num_samples", str(int(p.get("num_samples", 1))),
    "--device", "mps",
    "--output", str(workdir / stem),
]
seed = int(p.get("seed", -1))
if seed >= 0:
    cmd += ["--seed", str(seed)]
if p.get("no_postprocess"):
    cmd += ["--no-postprocess"]

t0 = time.time()
r = subprocess.run(cmd, cwd=str(ARDY), capture_output=True, text=True, env=env)
npz = sorted(workdir.rglob(f"{stem}*.npz"))
if r.returncode != 0 or not npz:
    print(r.stdout[-1500:])
    print(r.stderr[-2000:], file=sys.stderr)
    print("ERROR: ardy generation failed")
    sys.exit(1)
t_gen = time.time() - t0

outdir = Path(a.outdir)
slug = "".join(c if c.isalnum() else "_" for c in text[:40]).strip("_")
n = 0
for i, f in enumerate(npz):
    suffix = f"_{i:02d}" if len(npz) > 1 else ""
    npz_out = outdir / f"motion_{slug}{suffix}.npz"
    shutil.copy(f, npz_out)
    n += 1
    rb = subprocess.run(
        [sys.executable, str(ARDY2BVH), str(f), "-o", str(outdir / f"motion_{slug}{suffix}.bvh")],
        capture_output=True, text=True, env=env,
    )
    if rb.returncode == 0:
        n += 1
    else:
        print(f"WARN: bvh conversion failed for {f.name}: {rb.stderr[-500:]}", file=sys.stderr)

shutil.rmtree(workdir, ignore_errors=True)
if not any(outdir.glob("*.bvh")):
    print("ERROR: no BVH produced")
    sys.exit(1)
print(f"done: gen {t_gen:.0f}s, {n} files for '{text[:60]}'", flush=True)