import time
import gradio as gr
import os

import numpy as np
from PIL import Image
import math
from collections import defaultdict
import os
from huggingface_hub import HfApi


def greet(name):
    return "Hello " + name + "!"


def get_dataset_classes():
    # return {
    #     "airplane": [
    #         "airplane/0.glb",
    #         "airplane/1.glb",
    #         "airplane/2.glb",
    #         "airplane/3.glb",
    #         "airplane/4.glb",
    #         "airplane/5.glb",
    #         "airplane/6.glb",
    #         "airplane/7.glb",
    #     ],
    #     "bench": [
    #         "bench/0.glb",
    #         "bench/1.glb",
    #         "bench/10.glb",
    #         "bench/11.glb",
    #         "bench/12.glb",
    #         "bench/13.glb",
    #         "bench/14.glb",
    #         "bench/2.glb",
    #         "bench/3.glb",
    #         "bench/4.glb",
    #         "bench/5.glb",
    #         "bench/6.glb",
    #         "bench/7.glb",
    #         "bench/8.glb",
    #         "bench/9.glb",
    #     ],
    #     "camera": [
    #         "camera/0.glb",
    #         "camera/1.glb",
    #         "camera/2.glb",
    #         "camera/3.glb",
    #         "camera/4.glb",
    #         "camera/5.glb",
    #         "camera/6.glb",
    #         "camera/7.glb",
    #     ],
    #     "dishwasher": [
    #         "dishwasher/0.glb",
    #         "dishwasher/1.glb",
    #         "dishwasher/10.glb",
    #         "dishwasher/11.glb",
    #         "dishwasher/2.glb",
    #         "dishwasher/3.glb",
    #         "dishwasher/4.glb",
    #         "dishwasher/5.glb",
    #         "dishwasher/6.glb",
    #         "dishwasher/7.glb",
    #         "dishwasher/8.glb",
    #         "dishwasher/9.glb",
    #     ],
    #     "jar": [
    #         "jar/0.glb",
    #         "jar/1.glb",
    #         "jar/2.glb",
    #         "jar/3.glb",
    #         "jar/4.glb",
    #         "jar/5.glb",
    #         "jar/6.glb",
    #         "jar/7.glb",
    #         "jar/8.glb",
    #     ],
    #     "motorcycle": [
    #         "motorcycle/0.glb",
    #         "motorcycle/1.glb",
    #         "motorcycle/10.glb",
    #         "motorcycle/2.glb",
    #         "motorcycle/3.glb",
    #         "motorcycle/4.glb",
    #         "motorcycle/5.glb",
    #         "motorcycle/6.glb",
    #         "motorcycle/7.glb",
    #         "motorcycle/8.glb",
    #         "motorcycle/9.glb",
    #     ],
    #     "printer": [
    #         "printer/0.glb",
    #         "printer/1.glb",
    #         "printer/10.glb",
    #         "printer/11.glb",
    #         "printer/2.glb",
    #         "printer/3.glb",
    #         "printer/4.glb",
    #         "printer/5.glb",
    #         "printer/6.glb",
    #         "printer/7.glb",
    #         "printer/8.glb",
    #         "printer/9.glb",
    #     ],
    #     "sofa": [
    #         "sofa/0.glb",
    #         "sofa/1.glb",
    #         "sofa/10.glb",
    #         "sofa/11.glb",
    #         "sofa/12.glb",
    #         "sofa/2.glb",
    #         "sofa/3.glb",
    #         "sofa/4.glb",
    #         "sofa/5.glb",
    #         "sofa/6.glb",
    #         "sofa/7.glb",
    #         "sofa/8.glb",
    #         "sofa/9.glb",
    #     ],
    #     "washer": [
    #         "washer/0.glb",
    #         "washer/1.glb",
    #         "washer/10.glb",
    #         "washer/11.glb",
    #         "washer/2.glb",
    #         "washer/3.glb",
    #         "washer/4.glb",
    #         "washer/5.glb",
    #         "washer/6.glb",
    #         "washer/7.glb",
    #         "washer/8.glb",
    #         "washer/9.glb",
    #     ],
    # }

    hf_api = HfApi()
    info = hf_api.dataset_info(repo_id="ShapeNet/Example")

    dataset_classes = defaultdict(list)

    for file in info.siblings:

        if ".glb" in file.rfilename:
            class_name = file.rfilename.split("/")[0]
            dataset_classes[class_name].append(file.rfilename)
    print(dataset_classes)

    return dataset_classes
dataset_dict = get_dataset_classes()
dataset_classes = list(dataset_dict.keys())
default_models = dataset_dict[dataset_classes[0]]


def load_mesh(mesh_file_name):
    return mesh_file_name, mesh_file_name

def update(model_name):
    # wget the glb file from the datasets repo
    print(model_name)
    return f"{model_name}/0.glb"

def update_model_list(choice_class):
    print(f"inp1 changed {choice_class}")
    return {"choices":dataset_dict[choice_class]}


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inp = gr.Dropdown(choices=dataset_classes, interactive=True, label="3D Model Class", value=dataset_classes[0])
            out1 = gr.Dropdown(choices=default_models, interactive=True, label="3D Model", value=default_models[0])
            
            
        out2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model")
    
    inp.change(fn=update, inputs=inp, outputs=out2)
    # with gr.Row():
    #     btn = gr.Button("Load model")
    #     btn.click(fn=update, inputs=inp, outputs=out2)


# demo = gr.Interface(
#     fn=load_mesh,
#     inputs=gr.Model3D(),
#     outputs=[
#         gr.Model3D(
#             clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model"),
#         gr.File(label="Download 3D Model")
#     ],
#     examples=[
#     ],


demo.launch()