import os
import math
import shutil
import gradio as gr
from PIL import Image, ImageSequence
from moviepy.editor import VideoFileClip

TMP_DIR = "./__pycache__"


def clean_cache(tmp_dir=TMP_DIR):
    if os.path.exists(tmp_dir):
        shutil.rmtree(tmp_dir)

    os.mkdir(tmp_dir)
    return f"{tmp_dir}/input.gif"


def get_frame_duration(gif: Image):
    duration = gif.info.get("duration", 100)
    return [
        frame.info.get("duration", duration) for frame in ImageSequence.Iterator(gif)
    ]


def resize_gif(
    target_width: int,
    target_height: int,
    input_path=f"{TMP_DIR}/input.gif",
    output_path=f"{TMP_DIR}/output.gif",
):
    # Open the GIF image
    gif = Image.open(input_path)
    # Create a list to hold the modified frames
    modified_frames = []
    # Loop through each frame of the GIF
    for frame in ImageSequence.Iterator(gif):
        # Resize the frame
        resized_frame = frame.resize((target_width, target_height), Image.LANCZOS)
        # Append the resized frame to the list
        modified_frames.append(resized_frame)

    frame_durations = get_frame_duration(gif)
    modified_frames[0].save(
        output_path,
        format="GIF",
        append_images=modified_frames[1:],
        save_all=True,
        duration=frame_durations,
        loop=0,
    )

    return output_path


def infer(video_path: str, speed: float):
    target_w = 640
    gif_path = clean_cache()
    try:
        with VideoFileClip(video_path, audio_fps=16000) as clip:
            if clip.duration > 5:
                raise ValueError("The uploaded video is too long")

            # clip.write_gif(gif_path, fps=12, progress_bar=True)
            clip.speedx(speed).to_gif(gif_path, fps=12)
            w, h = clip.size

        target_h = math.ceil(target_w * h / w)
        return os.path.basename(video_path), resize_gif(target_w, target_h)

    except Exception as e:
        return f"{e}", None


if __name__ == "__main__":
    gr.Interface(
        fn=infer,
        inputs=[
            gr.Video(label="Upload video"),
            gr.Slider(
                label="Speed",
                minimum=0.5,
                maximum=2.0,
                step=0.25,
                value=1.0,
            ),
        ],
        outputs=[
            gr.Textbox(label="Filename", show_copy_button=True),
            gr.Image(label="Download GIF", type="filepath"),
        ],
        title="Please make sure the video is completely uploaded before clicking Submit, you can crop it online first if the video size is >5s",
        flagging_mode="never",
    ).launch()