import gradio as gr
import pandas as pd 
import requests
import json
from datasets import load_dataset

#dataset = load_dataset("nlphuji/flickr30k", split="test[10:20]")
#print (dataset)

#headers = {"Authorization": f"Bearer {API_TOKEN}"}
def query1(fetch_url):
    if fetch_url=="":
        fetch_url = "nlphuji/flickr30k"    
    API_URL1 = f"https://datasets-server.huggingface.co/splits?dataset={fetch_url}"
    response = requests.get(API_URL1)
    json_obj=response.json()

    config_box = []
    split_box = []
    for obj in json_obj:
        print (obj)
        if obj == "splits":
            try:
                for split in json_obj[obj]:
                    print (f'split: {split}')
                    out_config = split['config']
                    out_split = split['split']
                    print (out_config)
                    print (out_split)
                    config_box.append(out_config)
                    split_box.append(out_split)
            except Exception as e:
                print (e)
                pass 
    #df = pd.DataFrame.from_dict(json_obj)        
    return json_obj,gr.update(choices=[c for c in config_box]),gr.update(choices=[s for s in split_box])


def query2(fetch_url,config="TEST",split="test",offset=0,length=10):
    if fetch_url=="":
        fetch_url = "nlphuji/flickr30k"
    #offset=0
    #length=10
    API_URL2 = f"https://datasets-server.huggingface.co/rows?dataset={fetch_url}&config={config}&split={split}&offset={offset}&length={length}"
    response = requests.get(API_URL2)
    dictionary=response.json()
    #print (len(dictionary))
    lvl_1=[]
    
    for key in dictionary:
        lvl_1.append(key)
        print (key)
    return dictionary,gr.update(choices=[l for l in lvl_1])

def find_fn(lvl_1,lvl_2,fetch_url,config="TEST",split="test",offset=0,length=10):
    #print(out_json['rows'])
    #print (inp)
    img_list=[]
    out_json,_=query2(fetch_url,config,split,offset=10,length=20)
    lvl_1=lvl_1.strip("[]")

    lvl_2=lvl_2.strip('""').strip("''")
    #print(inp)
    cnt_lvl = ""
    if lvl_2 != "":
        ea_lvl_len = len(lvl_2)
        print (ea_lvl_len)
        for ea_lvl in lvl_2:
            cnt_lvl = f'{cnt_lvl}[{ea_lvl}]'
        for ea in out_json[f"{lvl_1}"]:
            cnt_lvl=cnt_lvl.strip("[]")
            #img_ea = ea['row']['image']['src']
            img_list.append(ea[f"{cnt_lvl}"])
    if lvl_2 == "":
        img_list.append(out_json[lvl_1])
    return img_list
def upd_drop(lvl_1=None,lvl_2=None,lvl_3=None,lvl_4=None,lvl_5=None,fetch_url="",config="TEST",split="test",offset=0,length=10):
    out_json,_=query2(fetch_url,config,split,offset=10,length=20)
    box_1=[]
    box_2=[]
    box_3=[]
    box_4=[]
    box_5=[]
    out1=None 
    out2=None 
    out3=None 
    out4=None 
    if lvl_1 == "rows":
        lvl_src= out_json[lvl_1][0]
    else:
        lvl_src=out_json[lvl_1]
    print (f"lvl_src:: {lvl_src}")
    if lvl_1 !="" and lvl_1 != None:
        for ea in (lvl_src):
            box_1.append(ea)    
        out= gr.update(choices = [m for m in box_1])
    if lvl_2 !="" and lvl_2 != None:
        for ea in (lvl_src[lvl_2]):
            print (ea)
            box_2.append(ea)    
        out= gr.update(choices = [m for m in box_2])
    else:
        pass
    if lvl_3 !="" and lvl_3 != None:
        for ea in (lvl_src):
            box_3.append(ea[lvl_2][lvl_3])    
        out= gr.update(choices = [m for m in box_3])
    else:
        pass
    if lvl_4 !="" and lvl_4 != None:
        for ea in (lvl_src):
            box_4.append(ea[lvl_2][lvl_3][lvl_4])    
        out= gr.update(choices = [m for m in box_4])  
    else:
        pass
    return out
def find_items(lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,fetch_url,config="TEST",split="test",offset=0,length=10):
    out_json,_=query2(fetch_url,config,split,offset=10,length=20)
    box_1=[]
    box_2=[]
    box_3=[]
    box_4=[]
    box_5=[]
    if lvl_1 !="":
        for ea in (out_json[lvl_1]):
            box_1.append(ea)    
        out = box_1
    if lvl_2 !="":
        for ea in (out_json[lvl_1]):
            box_2.append(ea[lvl_2])    
        out = box_2
    if lvl_3 !="":
        for ea in (out_json[lvl_1]):
            box_3.append(ea[lvl_2][lvl_3])    
        out = box_3  
    if lvl_4 !="":
        for ea in (out_json[lvl_1]):
            box_4.append(ea[lvl_2][lvl_3][lvl_4])    
        out = box_4  
    if lvl_5 !="":
        for ea in (out_json[lvl_1]):
            box_5.append(ea[lvl_2][lvl_3][lvl_4][lvl_5])    
        out = box_5          
    return out
def img_fn(inp):
    img_box=[]
    for img in inp:
        out = Image.open(inp)
        img_box.append(out)
    return img_box
with gr.Blocks() as app:
    with gr.Box():
        with gr.Row():
            with gr.Column(scale=3):
                with gr.Row():
                    data_set_url=gr.Textbox(label="Dataset (repo/name)")
                    fetch_btn=gr.Button()
                with gr.Row():
                    config_set=gr.Dropdown(label="Config")
                    split_set=gr.Dropdown(label="Split")
                    #config_drop = gr.Dropdown(label="Config/Split", choices=[])
                    view_data=gr.Button(label="View")
                with gr.Row():
                    lvl_1=gr.Dropdown(label="lvl_1")
                    lvl_2=gr.Dropdown(label="lvl_2")
                    lvl_3=gr.Dropdown(label="lvl_3")
                    lvl_4=gr.Dropdown(label="lvl_4")
                    lvl_5=gr.Dropdown(label="lvl_5")
                    
                    find_btn=gr.Button("Search")
                with gr.Tab("JSON"):
                    out_find = gr.JSON()
                with gr.Tab("Gallery"):
                    gal_btn=gr.Button("Load Images")
                    out_gal = gr.Gallery()

            with gr.Column(scale=1):
                out_json = gr.JSON()

    lvl_1.change(upd_drop,[lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,data_set_url,config_set,split_set],[lvl_2])
    lvl_2.change(upd_drop,[lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,data_set_url,config_set,split_set],[lvl_3])
    lvl_3.change(upd_drop,[lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,data_set_url,config_set,split_set],[lvl_4])
    lvl_4.change(upd_drop,[lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,data_set_url,config_set,split_set],[lvl_5])

    
    find_btn.click(find_items,[lvl_1,lvl_2,lvl_3,lvl_4,lvl_5,data_set_url,config_set,split_set],out_find)
    view_data.click(query2,[data_set_url,config_set,split_set],[out_find,lvl_1])
    gal_btn.click(img_fn,out_find,out_gal)
    
    
    #find_btn.click(find_fn,[lvl_1,lvl_2,data_set_url,config_set,split_set],[out_find])
    fetch_btn.click(query1,data_set_url,[out_json,config_set,split_set])
app.launch()