"""
Demonstrates integrating Rerun visualization with Gradio.

Provides example implementations of data streaming, keypoint annotation, and dynamic
visualization across multiple Gradio tabs using Rerun's recording and visualization capabilities.
"""

import math
import os
import tempfile
import time
import uuid

import cv2
import gradio as gr
import rerun as rr
import rerun.blueprint as rrb
from gradio_rerun import Rerun
from gradio_rerun.events import (
    SelectionChange,
    TimelineChange,
    TimeUpdate,
)
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from huggingface_hub import HfApi, HfFileSystem
from pathlib import Path

API   = HfApi()          # thin REST wrapper
FS    = HfFileSystem()   # fsspec-style filesystem               :contentReference[oaicite:0]{index=0}
REPO  = "pablovela5620/ego-dex-rrd"


SEQUENCE_NAMES = [
    "add_remove_lid",
    "arrange_topple_dominoes",
    "assemble_disassemble_furniture_bench_chair",
    "assemble_disassemble_furniture_bench_desk",
    "assemble_disassemble_furniture_bench_drawer",
    "assemble_disassemble_furniture_bench_lamp",
    "assemble_disassemble_furniture_bench_square_table",
    "assemble_disassemble_furniture_bench_stool",
    "assemble_disassemble_legos",
    "assemble_disassemble_soft_legos",
    "assemble_disassemble_structures",
    "assemble_disassemble_tiles",
    "assemble_jenga",
    "basic_fold",
    "basic_pick_place",
    "boil_serve_egg",
    "braid_unbraid",
    "build_unstack_lego",
    "charge_uncharge_airpods",
    "charge_uncharge_device",
    "clean_cups",
    "clean_surface",
    "clean_tableware",
    "clip_unclip_papers",
    "color",
    "crumple_flatten_paper",
    "deal_gather_cards",
    "declutter_desk",
    "dry_hands",
    "fidget_magnetic_spinner_rings",
    "flip_coin",
    "flip_pages",
    "fold_stack_unstack_unfold_cloths",
    "fold_unfold_paper_basic",
    "fold_unfold_paper_origami",
    "fry_bread",
    "fry_egg",
    "gather_roll_dice",
    "insert_dump_blocks",
    "insert_remove_airpods",
    "insert_remove_bagging",
    "insert_remove_bookshelf",
    "insert_remove_cups_from_rack",
    "insert_remove_drawer",
    "insert_remove_furniture_bench_cabinet",
    "insert_remove_furniture_bench_round_table",
    "insert_remove_plug_socket",
    "insert_remove_shirt_in_tube",
    "insert_remove_tennis_ball",
    "insert_remove_usb",
    "insert_remove_utensils",
    "knead_slime",
    "load_dispense_ice",
    "lock_unlock_key",
    "make_sandwich",
    "measure_objects",
    "open_close_insert_remove_box",
    "open_close_insert_remove_case",
    "open_close_insert_remove_tupperware",
    "paint_clean_brush",
    "peel_place_sticker",
    "pick_place_food",
    "pick_up_and_put_down_case_or_bag",
    "play_mancala",
    "play_piano",
    "play_reset_connect_four",
    "point_and_click_remote",
    "pour",
    "push_pop_toy",
    "put_away_set_up_board_game",
    "put_in_take_out_glasses",
    "put_toothpaste_on_toothbrush",
    "rake_smooth_zen_garden",
    "roll_ball",
    "scoop_dump_ice",
    "screw_unscrew_allen_fixture",
    "screw_unscrew_bottle_cap",
    "screw_unscrew_fingers_fixture",
    "set_up_clean_up_chessboard",
    "setup_cleanup_table",
    "sleeve_unsleeve_cards",
    "slot_batteries",
    "sort_beads",
    "stack",
    "stack_remove_jenga",
    "stack_unstack_bowls",
    "stack_unstack_cups",
    "stack_unstack_plates",
    "stack_unstack_tupperware",
    "staple_paper",
    "stock_unstock_fridge",
    "sweep_dustpan",
    "thread_unthread_bead_necklace",
    "throw_and_catch_ball",
    "throw_collect_objects",
    "tie_and_untie_shoelace",
    "tie_untie_rubberband",
    "type_keyboard",
    "use_chopsticks",
    "use_rubiks_cube",
    "vertical_pick_place",
    "wash_fruit",
    "wash_kitchen_dishes",
    "wash_put_away_dishes",
    "wipe_kitchen_surfaces",
    "wipe_screen",
    "wrap",
    "wrap_unwrap_food",
    "write",
    "zip_unzip_bag",
    "zip_unzip_case",
]

def show_dataset(task_name: str, episode_index: str):
    episode_index = f"{int(episode_index):05d}"
    url_str = f"https://huggingface.co/datasets/pablovela5620/ego-dex-rrd/resolve/main/{task_name}/{episode_index}.rrd"
    return url_str

def list_episodes(task:str) -> list[str]:
    """
    Return ["00000", "00001", ...] for the chosen task folder.
    """
    # fastest: one HTTP hit that returns the whole tree once
    files = API.list_repo_files(REPO, repo_type="dataset")
    return sorted(
        {Path(f).stem for f in files
         if f.startswith(f"{task}/") and f.endswith(".rrd")}
    )
default_task = SEQUENCE_NAMES[0]              # "add_remove_lid"
initial_eps  = list_episodes(default_task)    # ["00000", "00001", …]

with gr.Blocks() as demo:
    with gr.Tab("Hosted RRD"):
        with gr.Row():
            with gr.Column(scale=1):
                task_name = gr.Dropdown(
                    label="Task Name",
                    choices=SEQUENCE_NAMES,
                    value=default_task,                       
                )
                episode_index = gr.Dropdown(
                    label="Episode Index",
                    choices=initial_eps,                      
                    value=initial_eps[0] if initial_eps else None,

                )

                def _update_eps(t):                        # Gradio wants a fn
                    eps = list_episodes(t)
                    return gr.update(choices=eps, value=eps[0] if eps else None)

                task_name.change(_update_eps, inputs=task_name, outputs=episode_index)

                button = gr.Button("Show Dataset")
            with gr.Column(scale=4):
                viewer = Rerun(
                    streaming=True,
                    panel_states={
                        "time": "collapsed",
                        "blueprint": "hidden",
                        "selection": "hidden",
                    },
                )
        # choose_rrd.change(lambda x: x, inputs=[choose_rrd], outputs=[viewer])
        button.click(
            fn=show_dataset,
            inputs=[task_name, episode_index],
            outputs=[viewer]
        )




if __name__ == "__main__":
    demo.launch(ssr_mode=False)