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)

prompts = ds.load_dataset(
    "meta-ai-for-media-research/movie_gen_video_bench",
    split="test",
    streaming=True
)
prompts = [p["prompt"] for p in prompts][:num_videos]

with gr.Blocks() as demo:
    header = gr.Markdown(
        """
        # Goku MovieGen Video Bench Video Explorer
        Use this interface to view videos generated by [Goku](https://saiyan-world.github.io/goku/) using the promps from [Meta AI's MovieGen Video Bench](https://huggingface.co/datasets/meta-ai-for-media-research/movie_gen_video_bench). 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=prompts,
        value=prompts[0]
    )

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

    def decrement(current_prompt):
        current_id = prompts.index(current_prompt)
        return gr.Dropdown(
            label="Video",
            choices=prompts,
            value=prompts[max(0, current_id-1)]
        )

    def increment(current_prompt):
        current_id = prompts.index(current_prompt)
        return gr.Dropdown(
            label="Video",
            choices=prompts,
            value=prompts[min(num_videos-1, current_id+1)]
        )

    def random_video():
        return gr.Dropdown(
            label="Video",
            choices=prompts,
            value=prompts[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)