import gradio as gr
import pandas as pd
import requests

API_BASE = "https://datasets-server.huggingface.co"


def fetch(endpoint, params):
    try:
        response = requests.get(f"{API_BASE}/{endpoint}", params=params, timeout=30)
        response.raise_for_status()
        return response.json()
    except Exception as exc:
        return {"error": str(exc)}


def get_splits(dataset):
    dataset = (dataset or "").strip()
    if not dataset:
        return (
            pd.DataFrame(),
            gr.update(choices=[], value=None),
            "⚠️ Please enter a dataset name.",
        )

    data = fetch("splits", {"dataset": dataset})
    if "error" in data:
        return (
            pd.DataFrame(),
            gr.update(choices=[], value=None),
            f"❌ {data['error']}",
        )

    splits = data.get("splits", [])
    if not splits:
        return (
            pd.DataFrame(),
            gr.update(choices=[], value=None),
            "No splits found.",
        )

    df = pd.DataFrame(splits)
    columns = [c for c in ["dataset", "config", "split"] if c in df.columns]
    if columns:
        df = df[columns]

    choices = [f"{s.get('config', '')} :: {s.get('split', '')}" for s in splits]

    return (
        df,
        gr.update(choices=choices, value=choices[0] if choices else None),
        f"✅ Found {len(splits)} splits.",
    )


def get_rows(dataset, config_split, num_rows):
    dataset = (dataset or "").strip()
    if not dataset:
        return pd.DataFrame(), pd.DataFrame(), "⚠️ Please enter a dataset name."

    config = None
    split = None

    if config_split:
        if "::" in config_split:
            config, split = [x.strip() for x in config_split.split("::", 1)]
        else:
            split = config_split.strip()

    # Auto-fill missing config/split info using the splits endpoint
    if not config or not split:
        data = fetch("splits", {"dataset": dataset})
        if "error" in data:
            return pd.DataFrame(), pd.DataFrame(), f"❌ {data['error']}"

        splits = data.get("splits", [])
        if not splits:
            return pd.DataFrame(), pd.DataFrame(), "No splits found."

        if not split:
            if config:
                matching = [s for s in splits if s.get("config") == config]
                if not matching:
                    return pd.DataFrame(), pd.DataFrame(), f"❌ Config '{config}' not found."
                split = matching[0].get("split", "")
            else:
                config = splits[0].get("config", "")
                split = splits[0].get("split", "")
        else:
            for s in splits:
                if s.get("split") == split and (not config or s.get("config") == config):
                    config = s.get("config", "")
                    break
            else:
                return pd.DataFrame(), pd.DataFrame(), f"❌ Split '{split}' not found."

    data = fetch(
        "rows",
        {
            "dataset": dataset,
            "config": config,
            "split": split,
            "offset": 0,
            "length": int(num_rows),
        },
    )

    if "error" in data:
        return pd.DataFrame(), pd.DataFrame(), f"❌ {data['error']}"

    # Features / schema
    features = data.get("features", [])
    features_df = pd.DataFrame(features)
    if not features_df.empty and "type" in features_df.columns:
        features_df["type"] = features_df["type"].astype(str)

    # Rows
    rows = data.get("rows", [])
    if not rows:
        return features_df, pd.DataFrame(), "No rows returned."

    records = [r.get("row", {}) for r in rows]
    rows_df = pd.DataFrame(records)

    if rows and "row_idx" in rows[0]:
        rows_df.index = [r.get("row_idx", i) for i, r in enumerate(rows)]
        rows_df.index.name = "row_idx"

    return (
        features_df,
        rows_df,
        f"✅ Loaded {len(rows)} rows from `{config}` / `{split}`.",
    )


with gr.Blocks(title="HF Dataset Viewer") as demo:
    gr.Markdown(
        """
        # 🧐 Hugging Face Dataset Viewer

        Browse any public Hugging Face dataset through the
        [datasets-server API](https://datasets-server.huggingface.co).
        No dataset download needed.
        """
    )

    with gr.Row():
        dataset_text = gr.Textbox(
            label="Dataset name",
            value="imdb",
            placeholder="e.g. emotions, glue, squad",
        )
        splits_button = gr.Button("Get Splits")

    splits_table = gr.Dataframe(label="Splits", interactive=False)
    splits_status = gr.Markdown()

    with gr.Row():
        config_split_dropdown = gr.Dropdown(
            label="Config / Split",
            choices=[],
            interactive=True,
        )
        num_rows_slider = gr.Slider(
            minimum=1,
            maximum=100,
            value=10,
            step=1,
            label="Number of rows",
        )
        rows_button = gr.Button("Load Rows")

    with gr.Row():
        features_table = gr.Dataframe(label="Features", interactive=False)
        rows_table = gr.Dataframe(label="Rows", interactive=False)

    rows_status = gr.Markdown()

    splits_button.click(
        get_splits,
        inputs=dataset_text,
        outputs=[splits_table, config_split_dropdown, splits_status],
    )

    rows_button.click(
        get_rows,
        inputs=[dataset_text, config_split_dropdown, num_rows_slider],
        outputs=[features_table, rows_table, rows_status],
    )


if __name__ == "__main__":
    demo.launch()