import json
import gradio as gr
from huggingface_hub import hf_hub_download

DATASET_REPO = "atharva-manav/physicsbench-sampled"

TASKS = [
    "pb-hard-balance-hard-v1",
    "pb-hard-balance-medium-v1",
    "pb-hetbi-bowl-table-v1",
    "pb-hetbi-cracker-box-jaw-v1",
    "pb-hetbi-cracker-box-suction-v1",
    "pb-hetbi-cup-inversion-v1",
    "pb-hetbi-plate-shelf-v1",
    "pb-hobi-keyboard-typing-v1",
    "pb-hobi-plate-shelf-v1",
    "pb-pr-cup-extract-v1",
    "pb-pr-domino-select-v1",
    "pb-pr-domino-single-v1",
    "pb-pr-edge-slide-v1",
    "pb-pr-galileo-ramp-drop-v1",
    "pb-pr-mass-sort-v1",
    "pb-pr-mirror-pick-place-v1",
    "pb-pr-occluder-push-v1",
    "pb-pr-packing-v1",
    "pb-pr-recover-peg-insert-v1",
    "pb-pr-seesaw-balance-v1",
    "pb-pr-sheltered-grasp-v1",
    "pb-pr-slide-catch-v1",
    "pb-pr-slip-recovery-v1",
    "pb-pr-tool-retrieve-v1",
    "pb-puzzle-ball-cage-escape-v1",
    "pb-puzzle-gap-funnel-v1",
    "pb-puzzle-indirect-push-v1",
    "pb-puzzle-pick-place-ycb-v1",
    "pb-puzzle-ramp-sort-v1",
    "pb-puzzle-shape-stack-v1",
    "pb-puzzle-stack-collapse-recovery-v1",
]

_info_cache = {}


def get_info(task):
    if task not in _info_cache:
        try:
            path = hf_hub_download(
                repo_id=DATASET_REPO,
                filename=f"{task}/meta/info.json",
                repo_type="dataset",
            )
            with open(path) as f:
                _info_cache[task] = json.load(f)
        except Exception as e:
            return None
    return _info_cache[task]


def get_cameras(task):
    info = get_info(task)
    if not info:
        return []
    return [k for k, v in info["features"].items() if v.get("dtype") == "video"]


def on_task_change(task):
    cameras = get_cameras(task)
    info = get_info(task)
    n_ep = info.get("total_episodes", 5) if info else 5
    task_desc = info.get("task_id", task) if info else task
    return (
        gr.Dropdown(choices=cameras, value=cameras[0] if cameras else None, interactive=True),
        gr.Slider(minimum=0, maximum=n_ep - 1, step=1, value=0),
        task_desc,
    )


def load_video(task, episode_idx, camera):
    if not task or camera is None:
        return None, "Pick a task and camera then click Load."
    try:
        episode_idx = int(episode_idx)
        path = hf_hub_download(
            repo_id=DATASET_REPO,
            filename=f"{task}/videos/{camera}/chunk-000/file-{episode_idx:03d}.mp4",
            repo_type="dataset",
        )
        return path, ""
    except Exception as e:
        return None, f"Error loading video: {e}"


with gr.Blocks(title="PhysicsBench Visualizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# PhysicsBench Sampled Dataset Visualizer")
    gr.Markdown(
        "31 physics-manipulation tasks · 5 episodes each · multiple cameras  \n"
        f"Data: [{DATASET_REPO}](https://huggingface.co/datasets/{DATASET_REPO})"
    )

    with gr.Row():
        with gr.Column(scale=1):
            task_dd = gr.Dropdown(choices=TASKS, value=TASKS[0], label="Task", interactive=True)
            camera_dd = gr.Dropdown(choices=[], label="Camera", interactive=True)
            episode_slider = gr.Slider(minimum=0, maximum=4, step=1, value=0, label="Episode", interactive=True)
            task_desc = gr.Textbox(label="Task ID", interactive=False)
            load_btn = gr.Button("Load Video", variant="primary")
            err_box = gr.Textbox(label="", visible=True, interactive=False, show_label=False)

        with gr.Column(scale=2):
            video_out = gr.Video(label="Episode Video", autoplay=True)

    task_dd.change(on_task_change, inputs=[task_dd], outputs=[camera_dd, episode_slider, task_desc])
    load_btn.click(load_video, inputs=[task_dd, episode_slider, camera_dd], outputs=[video_out, err_box])
    demo.load(on_task_change, inputs=[task_dd], outputs=[camera_dd, episode_slider, task_desc])

demo.launch()