from typing import List, Dict
import httpx
import gradio as gr
import pandas as pd


async def get_splits(dataset_name: str) -> Dict[str, List[Dict]]:
    URL = f"https://datasets-server.huggingface.co/splits?dataset={dataset_name}"
    async with httpx.AsyncClient() as session:
        response = await session.get(URL)
        return response.json()


async def get_valid_datasets() -> Dict[str, List[str]]:
    URL = f"https://datasets-server.huggingface.co/valid"
    async with httpx.AsyncClient() as session:
        response = await session.get(URL)
        datasets = response.json()["valid"]
        return gr.Dropdown.update(choices=datasets, value="glue")


async def get_first_rows(dataset: str, config: str, split: str) -> Dict[str, Dict[str, List[Dict]]]:
    URL = f"https://datasets-server.huggingface.co/first-rows?dataset={dataset}&config={config}&split={split}"
    async with httpx.AsyncClient() as session:
        response = await session.get(URL)
        return response.json()

def get_df_from_rows(api_output):
    return pd.DataFrame([row["row"] for row in api_output["rows"]])


async def update_configs(dataset_name: str):
    splits = await get_splits(dataset_name)
    all_configs = sorted(set([s["config"] for s in splits["splits"]]))
    return (gr.Dropdown.update(choices=all_configs, value=all_configs[0]),
            splits)


async def update_splits(config_name: str, state: gr.State):
    splits_for_config = sorted(set([s["split"] for s in state["splits"] if s["config"] == config_name]))
    dataset_name = state["splits"][0]["dataset"]
    dataset = await update_dataset(splits_for_config[0], config_name, dataset_name)
    return (gr.Dropdown.update(choices=splits_for_config, value=splits_for_config[0]), dataset)


async def update_dataset(split_name: str, config_name: str, dataset_name: str):
    rows = await get_first_rows(dataset_name, config_name, split_name)
    df = get_df_from_rows(rows)
    return df


with gr.Blocks() as demo:
    splits_data = gr.State()
    with gr.Row():
        dataset_name = gr.Dropdown(label="Dataset")
        config = gr.Dropdown(label="Subset")
        split = gr.Dropdown(label="Split")
    with gr.Row():
        dataset = gr.DataFrame(wrap=True)
    demo.load(get_valid_datasets, inputs=None, outputs=[dataset_name])
    dataset_name.change(update_configs, inputs=[dataset_name], outputs=[config, splits_data])
    config.change(update_splits, inputs=[config, splits_data], outputs=[split, dataset])
    split.change(update_dataset, inputs=[split, config, dataset_name], outputs=[dataset])

demo.launch()