import subprocess
import uuid
import os
from pathlib import Path
import gradio as gr

# ---------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------
CKPT_PATH = "checkpoints/t2m_50step.pt"   # make sure this file exists
DEVICE    = "cpu"                         # free HF Spaces have no GPU
# ---------------------------------------------------------------------


def generate_motion(prompt: str) -> str:
    """
    Runs the MDM sampling script in a subprocess and returns the BVH
    file path so Gradio can hand it to the user.
    """
    out_file = Path("/tmp") / f"{uuid.uuid4().hex}.bvh"

    cmd = [
        "python",
        "-m",
        "motion_diffusion_model.sample.generate",
        "--model_path", str(CKPT_PATH),
        "--prompt",     prompt,
        "--output",     str(out_file),
        "--device",     DEVICE,
        "--num_steps",  "50",            # matches the checkpoint
    ]

    # --- make sure the local repo root is on PYTHONPATH so
    #     'utils.*' imports inside the script can be resolved
    env = os.environ.copy()
    root = Path(__file__).parent
    repo_inner = root / "motion_diffusion_model"
    env["PYTHONPATH"] = (
        f"{env.get('PYTHONPATH', '')}:{root}:{repo_inner}"
    )


    completed = subprocess.run(cmd, env=env, capture_output=True, text=True)
    if completed.returncode != 0:
        raise RuntimeError(f"Inference failed:\n{completed.stderr}")

    return str(out_file)


# -----------------------  Gradio UI  ----------------------------------
iface = gr.Interface(
    fn=generate_motion,
    inputs=gr.Textbox(
        lines=2,
        placeholder="e.g. a person walks forward and waves"
    ),
    outputs=gr.File(label="Download BVH"),
    title="Motion Diffusion Model – Text-to-Motion (50-step CPU demo)",
    description=(
        "Enter a natural-language prompt and receive a 3-D skeletal "
        "animation in BVH format."
    ),
)

if __name__ == "__main__":
    iface.launch()