import gradio as gr
from pathlib import Path
import subprocess
from inference_img import *  # Assuming necessary functions are exposed here

# Constants
OUTPUT_GIF = Path("demo/output.gif")
FRAME1_PATH = Path("demo/frame1.jpg")
FRAME2_PATH = Path("demo/frame2.jpg")

# Function to generate GIF
def generate_demo_gif(_1, _2):
    try:
        result = subprocess.run([
            "python", "inference_img.py",
            "--img", str(FRAME1_PATH), str(FRAME2_PATH),
            "--exp=4",
            "--model", "train_log/"
        ], capture_output=True, text=True)

        print("STDOUT:", result.stdout)
        print("STDERR:", result.stderr)

        if result.returncode == 0 and OUTPUT_GIF.exists():
            return str(OUTPUT_GIF), "✅ GIF generated successfully!"
        else:
            return None, "❌ GIF generation failed. Check logs or file paths."

    except Exception as e:
        print(f"Error: {e}")
        return None, f"❌ Subprocess error: {e}"

# Gradio UI
with gr.Blocks() as demo_ui:
    with gr.Tab("Demo"):
        gr.Markdown("### Preview frames and generate animation")
        with gr.Row():
            img1_preview = gr.Image(value=str(FRAME1_PATH), label="Frame 1", interactive=False)
            img2_preview = gr.Image(value=str(FRAME2_PATH), label="Frame 2", interactive=False)
        generate_btn = gr.Button("Generate GIF")
        output_gif = gr.Image(label="Animated GIF")
        status_msg = gr.Markdown()

        generate_btn.click(
            fn=generate_demo_gif,
            inputs=[img1_preview, img2_preview],  # Non-interactive but passed to match fn signature
            outputs=[output_gif, status_msg]
        )

demo_ui.launch()