import gradio as gr
import subprocess
import os
from pydub import AudioSegment
import tempfile

# Set OpenGL platform
os.environ["PYOPENGL_PLATFORM"] = "osmesa"

# Paths to TalkShow demo script and configuration
DEMO_SCRIPT_PATH = "scripts/demo.py"
CONFIG_FILE = "./config/LS3DCG.json"
BODY_MODEL_PATH = "experiments/2022-10-19-smplx_S2G-LS3DCG/ckpt-99.pth"
OUTPUT_MP4_PATH = "./output/demo_output.mp4"

# Ensure the output directory exists
os.makedirs(os.path.dirname(OUTPUT_MP4_PATH), exist_ok=True)

def convert_audio_to_wav(input_audio):
    # Use a temporary file to save the .wav output
    temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    temp_wav_path = temp_wav.name

    # Convert the input audio to .wav format
    audio = AudioSegment.from_file(input_audio)
    audio.export(temp_wav_path, format="wav")
    return temp_wav_path

def run_demo_and_get_video(input_audio):
    # Convert the input audio to a temporary .wav file
    wav_path = convert_audio_to_wav(input_audio)

    try:
        # Run the demo script with the specified arguments
        command = [
            "python", DEMO_SCRIPT_PATH,
            "--config_file", CONFIG_FILE,
            "--infer",
            "--audio_file", wav_path,
            "--body_model_name", "s2g_LS3DCG",
            "--body_model_path", BODY_MODEL_PATH,
            "--id", "0"
        ]
        subprocess.run(command, check=True)

        # Check if the output file exists
        if os.path.exists(OUTPUT_MP4_PATH):
            return OUTPUT_MP4_PATH
        else:
            return "Error: Output video not generated."
    except subprocess.CalledProcessError as e:
        return f"Error: {str(e)}"
    finally:
        # Clean up the temporary .wav file after processing
        os.remove(wav_path)

# Define the Gradio interface
interface = gr.Interface(
    fn=run_demo_and_get_video,
    inputs=gr.Audio(source="upload", type="filepath"),
    outputs=gr.Video(type="file"),
    title="TalkSHOW Demo",
    description="Generate a video from audio input using TalkSHOW model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()