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()

    drop_box = []
    for obj in json_obj:
        print (obj)
        out_config = obj[0]
        out_split = obj[1]
        print (out_config)
        print (out_split)
        drop_box.append(obj[0])
    #df = pd.DataFrame.from_dict(json_obj)        
    return json_obj


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()
    return dictionary

def find_fn(inp,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)
    
    inp=inp.strip("[]")
    print(inp)
    for ea in out_json[f"{inp}"]:
        #img_ea = ea['row']['image']['src']
        img_list.append(ea)
    
    return img_list
with gr.Blocks() as app:
    with gr.Row():
        data_set_url=gr.Textbox(label="Dataset (repo/name)")
        fetch_btn=gr.Button()
    with gr.Row():
        config_set=gr.Textbox(label="Config")
        split_set=gr.Textbox(label="Split")
        #config_drop = gr.Dropdown(label="Config/Split", choices=[])
        view_data=gr.Button(label="View")
    with gr.Row():
        find_string=gr.Textbox()
        find_btn=gr.Button("Search")
    with gr.Row():
        with gr.Column(scale=1):
            out_json = gr.JSON()
        with gr.Column(scale=3):
            out_find = gr.JSON()
            out_gal = gr.Gallery()

    view_data.click(query2,[data_set_url,config_set,split_set],[out_find])
    find_btn.click(find_fn,[find_string,data_set_url,config_set,split_set],[out_find])
    fetch_btn.click(query1,data_set_url,out_json)
app.launch()