import math
import os
import tempfile
import time
import uuid


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/spaces/DrChamyoung/DexDatasetRefiner/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)