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:
    out_dir = Path("/tmp") / f"mdm_{uuid.uuid4().hex}"
    out_dir.mkdir(parents=True, exist_ok=True)
    
    cmd = [
        "python", "-m", "motion_diffusion_model.sample.generate",
        "--model_path", str(CKPT_PATH),
        "--text_prompt", prompt,
        "--cuda", "-1",
        "--guidance_param", "7.5",
        "--output_dir", str(out_dir),
        "--num_samples", "1",
        "--unconstrained",
        "--inference_only"  # Add this flag (you'll need to implement it)
    ]
    
    env = os.environ.copy()
    root = Path(__file__).parent
    repo = root / "motion_diffusion_model"
    env["PYTHONPATH"] = f"{env.get('PYTHONPATH','')}:{root}:{repo}"
    
    try:
        subprocess.run(cmd, env=env, check=True)
        bvh = next(out_dir.rglob("*.bvh"))
        return str(bvh)
    except subprocess.CalledProcessError as e:
        return f"Error generating motion: {str(e)}"
    except StopIteration:
        return "No BVH file was generated."

# -----------------------  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()