import os
import random
import gradio as gr
import datasets as ds

bench = ds.load_dataset(
    "saiyan-world/Goku-MovieGenBench",
    split="train"
)
bench = bench.cast_column("video", ds.Video(decode=False))
num_videos = len(bench)

with gr.Blocks() as demo:
    header = gr.Markdown(
        """
        # Goku MovieGen Bench Video Explorer
        Use this interface to view videos generated by [Goku](https://saiyan-world.github.io/goku/). For more information, see [Goku: Flow Based Video Generate Models on arXiv](https://arxiv.org/abs/2502.04896).                                       """.strip()
    )
    video = gr.Video(
        value=bench[0]["video"]["path"],
        loop=True,
        autoplay=True,
    )

    dropdown = gr.Dropdown(
        label="Video",
        choices=range(num_videos),
        value=0
    )

    def set_video(new_id):
        return gr.Video(
            value=bench[new_id]["video"]["path"],
            loop=True,
            autoplay=True,
        )

    def decrement(current_id):
        return gr.Dropdown(
            label="Video",
            choices=range(num_videos),
            value=max(0, current_id-1)
        )

    def increment(current_id):
        return gr.Dropdown(
            label="Video",
            choices=range(num_videos),
            value=min(num_videos-1, current_id+1)
        )

    def random_video():
        return gr.Dropdown(
            label="Video",
            choices=range(num_videos),
            value=random.randint(0, num_videos-1)
        )

    with gr.Row():
        go_previous = gr.Button("Previous")
        go_previous.click(
            fn=decrement,
            inputs=dropdown,
            outputs=dropdown,
        )
        go_next = gr.Button("Next")
        go_next.click(
            fn=increment,
            inputs=dropdown,
            outputs=dropdown
        )
        go_random = gr.Button("Random")
        go_random.click(
            fn=random_video,
            outputs=dropdown,
        )

    dropdown.change(fn=set_video, inputs=dropdown, outputs=video)

paths = [
    os.path.dirname(bench[0]["video"]["path"]),
    os.path.expanduser("~/.cache/huggingface")
]

demo.launch(allowed_paths=paths)