import gradio as gr
import os
import shutil
import glob

# ------------------------------
# 1. LIA MODEL RUN FUNCTION
# ------------------------------
def run_lia_model(source_image_path, driving_video_path):
    if source_image_path is None or driving_video_path is None:
        return None

    # Clean names for LIA expected input paths
    source_clean = "source_image.jpg"
    driving_clean = "driving_video.mp4"

    # Remove previous files if they exist
    if os.path.exists(source_clean): os.remove(source_clean)
    if os.path.exists(driving_clean): os.remove(driving_clean)

    shutil.copy(source_image_path, source_clean)
    shutil.copy(driving_video_path, driving_clean)

    print("🚀 Running LIA model...")

    # Run the LIA script
    exit_code = os.system(
        f'python run_demo.py --model vox --source_path {source_clean} --driving_path {driving_clean}'
    )

    if exit_code != 0:
        print("❌ Error running run_demo.py")
        return None

    # Look for latest result video
    results = glob.glob('res/vox/*.mp4')
    if not results:
        print("❌ No output video found.")
        return None

    latest = max(results, key=os.path.getctime)
    print(f"🎉 Returning output: {latest}")
    return latest


# ------------------------------
# 2. GRADIO INTERFACE
# ------------------------------
interface = gr.Interface(
    fn=run_lia_model,
    inputs=[
        gr.Image(type="filepath", label="Source Image (Face Image)"),
        gr.Video(label="Driving Video (Motion Video)")
    ],
    outputs=gr.Video(label="Generated Animation"),
    title="LIA - Latent Image Animator",
    description="Upload a face image and driving video to generate an animation."
)

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