# app.py
import os
import uuid
import subprocess
from pathlib import Path
from typing import Tuple
from inference_img import *
import gradio as gr

# Demo images
FRAME1 = Path("demo/frame1.jpg")
FRAME2 = Path("demo/frame2.jpg")
DEMO = Path("demo/demo.gif")
MODEL =  Path("train_log/")

TMP_PREFIX = "demo/"
TARGET_DIR = f"{TMP_PREFIX}output/"
EXTENSION = "png"

def interpolate_image(img_a_path: str, img_b_path: str) -> Tuple[str, str]:
    """
    3-step pipeline:
    1) inference_img.py --img <A> <B> --exp=4  -> writes frames as /tmp/gradio/output/img%d.png
    2) ffmpeg palettegen
    3) ffmpeg paletteuse -> final GIF at a unique path
    Returns: (gif_filepath_for_preview, gif_filepath_for_download)
    """
    # Ensure temp dirs exist
    os.makedirs(TARGET_DIR, exist_ok=True)

    # Unique output path per run
    interpolated_path = f"{TARGET_DIR}{uuid.uuid4()}.{EXTENSION}"
    palette_path = f"{TMP_PREFIX}palette.png"

    try:

        #inference_img(img=[Path(img_a_path), Path(img_b_path)],exp=4)
        
        print("ok")

subprocess.run([
            "ffmpeg", "-y", "-r", "14", "-f", "image2",
            "-i", f"{TARGET_DIR}img%d.png",
            "-vf", "palettegen=stats_mode=single",
            palette_path
        ], check=True)



        
        # 1) Run interpolation (frames -> TARGET_DIR/img1.png ... imgN.png)
        #subprocess.run([
        #        "python3", "inference_img.py",
        #         "--img", str(img_a_path), str(img_b_path),
        #         "--exp", "4"
        #], check=True)

        print("ok")
        # 2) Generate palette
        subprocess.run([
            "ffmpeg", "-y", "-r", "14", "-f", "image2",
            "-i", f"{TARGET_DIR}img%d.png",
            "-vf", "palettegen=stats_mode=single",
            palette_path
        ], check=True)

        # 3) Use palette to produce final GIF
        #!ffmpeg -r 14 -f image2 -i output/img%d.png -vf "palettegen=stats_mode=single" palette.png
        subprocess.run([
            "ffmpeg", "-y", "-r", "14", "-f",
            "-i", f"{TARGET_DIR}img%d.png",
            "palettegen=stats_mode=single",
            palette_path
        ], check=True)

        return interpolated_image, interpolated_path

    except subprocess.CalledProcessError as e:
        raise gr.Error(f"Interpolation failed. {e}")

# Gradio UI
with gr.Blocks(title="RIFE Interpolation") as demo:
    with gr.Tab("Demo"):
        gr.Markdown("### Demo: Preloaded images")
        input_imageA = gr.Image(type="filepath", value=str(FRAME1), label="Image A")
        input_imageB = gr.Image(type="filepath", value=str(FRAME2), label="Image B")
        output_image = gr.Image(type="filepath", value=str(DEMO), label="Demo")
        g_btn = gr.Button("Interpolate")

        output_image = gr.Image(type="filepath", label="Interpolated GIF")
        enhance_image_path = gr.Textbox(label="Output path", interactive=False)

        g_btn.click(
            fn=interpolate_image,
            inputs=[input_imageA, input_imageB]
        )

    with gr.Tab("Upload your images"):
        gr.Markdown("### Upload any two images")
        input_imageA = gr.Image(type="filepath", label="Image A")
        input_imageB = gr.Image(type="filepath", label="Image B")
        user_btn = gr.Button("Interpolate")
        user_out = gr.Image(type="filepath", value=str(TARGET_DIR), label="Interpolated GIF")
        user_path = gr.Textbox(label="Output path", value=str(TARGET_DIR), interactive=False)
        user_btn.click(
            fn=interpolate_image,
            inputs=[input_imageA, input_imageB],
            outputs=[user_out, user_path],
        )

demo.launch()