app.py

"""
HF Dataset Explorer — view ANY Hugging Face dataset, even when the
official dataset-viewer is unavailable / stuck / not-yet-built.

Run:  python app.py         (add --share for a public link)
"""

from __future__ import annotations

import argparse
import json
import os
import tempfile
import traceback

import gradio as gr
import pandas as pd

import backend as B

STATE_DEFAULT = {"dataset": "", "config": "", "split": "", "df": None,
                 "offset": 0, "total": None, "token": ""}

EXAMPLES = [
    "stanfordnlp/imdb",
    "openai/gsm8k",
    "cais/mmlu",
    "HuggingFaceFW/fineweb",
    "lmms-lab/COCO-Caption",
]


# --------------------------------------------------------------------------- #
# tab 1: connect
# --------------------------------------------------------------------------- #
def connect(raw_id, token):
    try:
        ds = B.normalize_id(raw_id)
    except B.LoadError as e:
        return (gr.update(), gr.update(), f"### ❌ {e}", "", gr.update(choices=[]), {})

    lines = [f"## `{ds}`", ""]
    try:
        info = B.repo_info(ds, token)
        lines += [
            f"- **Downloads (30d):** {info.get('downloads', '?'):,}  ·  **Likes:** {info.get('likes', '?')}",
            f"- **Modified:** {str(info.get('lastModified', ''))[:19]}",
            f"- **Tags:** {', '.join(info.get('tags', [])[:14]) or '—'}",
            f"- **Gated:** {info.get('gated', False)}  ·  **Private:** {info.get('private', False)}",
        ]
    except Exception as e:  # noqa: BLE001
        lines.append(f"> ⚠️ repo metadata unavailable: {e}")

    ok, detail = B.viewer_status(ds, token)
    lines += ["",
              f"- **Official viewer:** {'✅ ready' if ok else '⛔ NOT ready'} — `{detail}`",
              f"  {'' if ok else '→ no problem, this app falls back to parquet / streaming / raw files.'}"]

    try:
        mapping = B.discover_splits(ds, token)
    except Exception as e:  # noqa: BLE001
        mapping = {"default": ["train"]}
        lines.append(f"> ⚠️ split discovery fell back to guesses: {e}")

    configs = list(mapping)
    first_cfg = configs[0]
    splits = mapping[first_cfg]
    lines.append(f"- **Configs:** {len(configs)} · **Splits in `{first_cfg}`:** {', '.join(splits)}")

    # file tree
    try:
        files = B.list_files(ds, token)
        tab = pd.DataFrame([{"path": f["path"],
                             "size": f.get("size") or (f.get("lfs") or {}).get("size") or 0}
                            for f in files if f.get("type") == "file"])
        if not tab.empty:
            tab["size"] = tab["size"].map(lambda n: f"{n/1e6:,.2f} MB" if n else "—")
        tabular = [f["path"] for f in files
                   if f.get("type") == "file" and f["path"].lower().endswith(B.TABULAR_EXT)]
    except Exception:
        tab, tabular = pd.DataFrame(), []

    return (gr.update(choices=configs, value=first_cfg),
            gr.update(choices=splits, value=splits[0] if splits else None),
            "\n".join(lines),
            B.readme(ds, token),
            gr.update(choices=tabular[:500], value=tabular[0] if tabular else None),
            {"dataset": ds, "mapping": mapping, "files": tab.to_dict("records")})


def on_config_change(cfg, meta):
    splits = (meta or {}).get("mapping", {}).get(cfg, ["train"])
    return gr.update(choices=splits, value=splits[0] if splits else None)


def file_table(meta):
    recs = (meta or {}).get("files", [])
    return pd.DataFrame(recs) if recs else pd.DataFrame({"path": [], "size": []})


# --------------------------------------------------------------------------- #
# tab 2: rows
# --------------------------------------------------------------------------- #
def _render(res: B.LoadResult, state, maxlen, cols_pick=None):
    df = res.df
    state = dict(state or {})
    state.update(df=df, offset=res.offset, total=res.total_rows,
                 strategy=res.strategy)

    media = B.detect_media_columns(df)
    total = f"{res.total_rows:,}" if res.total_rows else "unknown"
    banner = (f"**Strategy:** `{res.strategy}` · **rows {res.offset:,}–{res.offset+len(df)-1:,}** "
              f"of {total} · **{len(df.columns)} columns**"
              + (f"\n\n_{' · '.join(res.notes)}_" if res.notes else "")
              + (f"\n\n🖼️ image cols: `{media['image']}`" if media["image"] else "")
              + (f" 🔊 audio cols: `{media['audio']}`" if media["audio"] else ""))

    keep = [c for c in (cols_pick or []) if c in df.columns]
    show = df[keep] if keep else df
    return (B.display_frame(show, maxlen), banner, state,
            gr.update(choices=[str(c) for c in df.columns],
                      value=[str(c) for c in df.columns]),
            gr.update(maximum=max(1, len(df) - 1), value=0))


def load_rows(dataset_meta, cfg, split, offset, limit, token, strategy, maxlen, state):
    ds = (dataset_meta or {}).get("dataset", "")
    if not ds:
        return (pd.DataFrame(), "### ❌ Connect to a dataset first (tab 1).", state,
                gr.update(), gr.update())
    try:
        res = B.load_rows(ds, cfg, split, int(offset), int(limit), token, strategy)
    except Exception as e:  # noqa: BLE001
        return (pd.DataFrame(), f"### ❌ Could not load rows\n```\n{e}\n```", state,
                gr.update(), gr.update())
    state = dict(state or {}); state["dataset"] = ds
    return _render(res, state, int(maxlen))


def page(delta, dataset_meta, cfg, split, offset, limit, token, strategy, maxlen, state):
    new_off = max(0, int(offset) + delta * int(limit))
    out = load_rows(dataset_meta, cfg, split, new_off, limit, token, strategy, maxlen, state)
    return (*out, new_off)


def _state_df(state) -> pd.DataFrame:
    df = (state or {}).get("df")
    return df if isinstance(df, pd.DataFrame) else pd.DataFrame()


def apply_filter(query, regex_col, regex, state, maxlen):
    df = _state_df(state)
    if df.empty:
        return pd.DataFrame(), "No data loaded."
    note = []
    try:
        if query.strip():
            df = df.query(query, engine="python")
            note.append(f"query `{query}`")
        if regex.strip() and regex_col:
            df = df[df[regex_col].astype(str).str.contains(regex, case=False, regex=True, na=False)]
            note.append(f"regex `{regex}` on `{regex_col}`")
    except Exception as e:  # noqa: BLE001
        return pd.DataFrame(), f"❌ filter error: `{e}`"
    return B.display_frame(df, int(maxlen)), f"**{len(df):,} matching rows** ({' + '.join(note) or 'no filter'})"


def pick_columns(cols, state, maxlen):
    df = _state_df(state)
    if df.empty:
        return pd.DataFrame()
    keep = [c for c in df.columns if str(c) in set(cols or [])]
    return B.display_frame(df[keep] if keep else df, int(maxlen))


# --------------------------------------------------------------------------- #
# tab 3: record inspector
# --------------------------------------------------------------------------- #
def inspect(idx, state):
    df = _state_df(state)
    if df.empty:
        return {}, "No data loaded.", None, None
    i = max(0, min(int(idx), len(df) - 1))
    row = df.iloc[i]
    payload = {str(k): B.json_safe(v) for k, v in row.items()}

    md = [f"### Row {i} (absolute #{(state or {}).get('offset', 0) + i})"]
    for k, v in row.items():
        s = B.stringify(v, 4000)
        md.append(f"**{k}**\n\n```\n{s}\n```")

    img = aud = None
    media = B.detect_media_columns(df)
    if media["image"]:
        img = B.media_to_display(row[media["image"][0]])
    if media["audio"]:
        aud = B.media_to_display(row[media["audio"][0]])
    return payload, "\n\n".join(md), img, aud


# --------------------------------------------------------------------------- #
# tab 4: stats
# --------------------------------------------------------------------------- #
def profile(state):
    df = _state_df(state)
    if df.empty:
        return pd.DataFrame(), pd.DataFrame(), "No data loaded."
    rows = []
    for c in df.columns:
        s = df[c]
        flat = s.map(lambda v: B.stringify(v, 10_000))
        rows.append({
            "column": str(c),
            "dtype": str(s.dtype),
            "nulls": int(s.isna().sum()),
            "null %": round(100 * s.isna().mean(), 2),
            "unique": int(flat.nunique()),
            "mean len": round(flat.str.len().mean(), 1),
            "max len": int(flat.str.len().max()),
            "example": B.stringify(s.dropna().iloc[0], 80) if s.notna().any() else "",
        })
    schema = pd.DataFrame(rows)

    num = df.select_dtypes("number")
    desc = num.describe().T.reset_index().rename(columns={"index": "column"}) if not num.empty \
        else pd.DataFrame({"info": ["no numeric columns in this page"]})

    # label-ish distribution
    md = ["### Value distributions (low-cardinality columns)"]
    for c in df.columns:
        flat = df[c].map(lambda v: B.stringify(v, 60))
        if 1 < flat.nunique() <= 25:
            vc = flat.value_counts()
            md.append(f"**{c}**")
            md += [f"- `{k}` — {v} ({100*v/len(flat):.1f}%)" for k, v in vc.items()]
    if len(md) == 1:
        md.append("_none found on this page._")
    return schema, desc, "\n".join(md)


# --------------------------------------------------------------------------- #
# tab 5: media gallery
# --------------------------------------------------------------------------- #
def gallery(state, col, n):
    df = _state_df(state)
    if df.empty:
        return [], "No data loaded."
    media = B.detect_media_columns(df)
    col = col or (media["image"][0] if media["image"] else None)
    if not col or col not in df.columns:
        return [], "No image-like column detected on this page."
    items = []
    for i, v in df[col].head(int(n)).items():
        d = B.media_to_display(v)
        if d is not None:
            items.append((d, f"#{i}"))
    return items, f"{len(items)} image(s) from `{col}`"


def media_choices(state):
    df = _state_df(state)
    if df.empty:
        return gr.update(choices=[])
    m = B.detect_media_columns(df)
    return gr.update(choices=m["image"] + m["audio"],
                     value=(m["image"] or m["audio"] or [None])[0])


# --------------------------------------------------------------------------- #
# tab 6: raw files
# --------------------------------------------------------------------------- #
def read_raw(dataset_meta, path, offset, limit, token, maxlen):
    ds = (dataset_meta or {}).get("dataset", "")
    if not ds or not path:
        return pd.DataFrame(), "Pick a file first."
    try:
        res = B.read_raw_file(ds, path, token, int(limit), int(offset))
    except Exception as e:  # noqa: BLE001
        return pd.DataFrame(), f"❌ {e}"
    return (B.display_frame(res.df, int(maxlen)),
            f"**{res.strategy}** · {len(res.df)} rows"
            + (f" of {res.total_rows:,}" if res.total_rows else ""))


# --------------------------------------------------------------------------- #
# tab 7: export
# --------------------------------------------------------------------------- #
def export(state, fmt):
    df = _state_df(state)
    if df.empty:
        return None, "Nothing to export."
    name = f"{(state or {}).get('dataset','hf')}".replace("/", "__") or "hf_export"
    d = tempfile.mkdtemp()
    path = os.path.join(d, f"{name}.{ {'CSV':'csv','JSON Lines':'jsonl','Parquet':'parquet'}[fmt] }")
    try:
        if fmt == "CSV":
            B.display_frame(df, 10**6).to_csv(path, index=False)
        elif fmt == "JSON Lines":
            with open(path, "w", encoding="utf-8") as fh:
                for _, r in df.iterrows():
                    fh.write(json.dumps({str(k): B.json_safe(v) for k, v in r.items()},
                                        ensure_ascii=False) + "\n")
        else:
            df.to_parquet(path, index=False)
    except Exception as e:  # noqa: BLE001
        return None, f"❌ {e}"
    return path, f"✅ exported {len(df):,} rows"


def code_snippet(meta, cfg, split, strategy):
    ds = (meta or {}).get("dataset", "your/dataset")
    return f'''```python
# Reproduce this page in your own code
from datasets import load_dataset

# 1) streaming — works even when the Hub viewer is not ready
ds = load_dataset("{ds}", {json.dumps(cfg) if cfg else "None"}, split="{split}", streaming=True)
for i, ex in zip(range(5), ds):
    print(ex)

# 2) direct parquet (random access, no dataset script executed)
import pandas as pd
df = pd.read_parquet(
    "hf://datasets/{ds}@refs/convert/parquet/{cfg or 'default'}/{split}/0000.parquet"
)
```'''


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
CSS = """
.small-note {font-size: 0.85rem; opacity: .75}
footer {display:none !important}
"""

with gr.Blocks(title="HF Dataset Explorer") as demo:
    meta = gr.State({})
    state = gr.State(dict(STATE_DEFAULT))

    gr.Markdown(
        "# 🤗 HF Dataset Explorer\n"
        "View **any** Hugging Face dataset — even when the official dataset viewer "
        "is not ready. Falls back automatically: `viewer API → parquet → streaming → raw files`."
    )

    with gr.Row():
        dataset_id = gr.Textbox(label="Dataset id or URL", value="stanfordnlp/imdb",
                                placeholder="user/name, or paste a huggingface.co URL", scale=4)
        token = gr.Textbox(label="HF token (optional, for private/gated)",
                           type="password", scale=2)
        connect_btn = gr.Button("Connect", variant="primary", scale=1)
    gr.Examples(EXAMPLES, inputs=dataset_id, label="Try")

    with gr.Row():
        config = gr.Dropdown(label="Config", choices=[], interactive=True, allow_custom_value=True)
        split = gr.Dropdown(label="Split", choices=[], interactive=True, allow_custom_value=True)
        strategy = gr.Dropdown(label="Loader", value="Auto (try everything)",
                               choices=["Auto (try everything)", *B.STRATEGIES])
        limit = gr.Slider(5, 500, 50, step=5, label="Rows per page")
        offset = gr.Number(0, label="Offset", precision=0)
        maxlen = gr.Slider(50, 2000, 300, step=50, label="Cell truncation (chars)")

    info_md = gr.Markdown()

    with gr.Tabs():
        with gr.Tab("📋 Rows"):
            with gr.Row():
                load_btn = gr.Button("Load rows", variant="primary")
                prev_btn = gr.Button("◀ Prev page")
                next_btn = gr.Button("Next page ▶")
            banner = gr.Markdown()
            cols = gr.Dropdown(label="Visible columns", multiselect=True, choices=[], interactive=True)
            table = gr.Dataframe(wrap=True, interactive=False, max_height=650)

        with gr.Tab("🔎 Filter & search"):
            gr.Markdown("Filters run on the **currently loaded page** (client-side, instant).")
            with gr.Row():
                query = gr.Textbox(label="pandas query", placeholder="label == 1 and len(text) > 500")
                regex_col = gr.Dropdown(label="Regex column", choices=[], allow_custom_value=True)
                regex = gr.Textbox(label="Regex / substring", placeholder="great|terrible")
            filter_btn = gr.Button("Apply filter", variant="primary")
            filter_note = gr.Markdown()
            ftable = gr.Dataframe(wrap=True, interactive=False, max_height=600)

        with gr.Tab("🔬 Record inspector"):
            row_idx = gr.Slider(0, 1, 0, step=1, label="Row on this page")
            with gr.Row():
                with gr.Column(scale=3):
                    rec_md = gr.Markdown()
                with gr.Column(scale=2):
                    rec_json = gr.JSON(label="Raw record")
                    rec_img = gr.Image(label="Image field", height=260)
                    rec_aud = gr.Audio(label="Audio field")

        with gr.Tab("📊 Stats"):
            stats_btn = gr.Button("Profile this page", variant="primary")
            schema_tbl = gr.Dataframe(label="Schema & quality", wrap=True, interactive=False)
            num_tbl = gr.Dataframe(label="Numeric describe()", interactive=False)
            dist_md = gr.Markdown()

        with gr.Tab("🖼️ Media"):
            with gr.Row():
                media_col = gr.Dropdown(label="Media column", choices=[], allow_custom_value=True)
                n_media = gr.Slider(1, 100, 24, step=1, label="How many")
                media_btn = gr.Button("Render", variant="primary")
            media_note = gr.Markdown()
            gal = gr.Gallery(columns=6, height=520, object_fit="contain")

        with gr.Tab("🗂️ Repo files"):
            gr.Markdown("Bypass everything and read a file straight out of the repo.")
            files_tbl = gr.Dataframe(interactive=False, max_height=280)
            with gr.Row():
                file_pick = gr.Dropdown(label="Tabular file", choices=[], allow_custom_value=True, scale=3)
                raw_off = gr.Number(0, label="Offset", precision=0)
                raw_lim = gr.Slider(5, 500, 50, step=5, label="Rows")
                raw_btn = gr.Button("Read file", variant="primary")
            raw_note = gr.Markdown()
            raw_tbl = gr.Dataframe(wrap=True, interactive=False, max_height=520)

        with gr.Tab("📄 README"):
            readme_md = gr.Markdown()

        with gr.Tab("⬇️ Export & code"):
            with gr.Row():
                fmt = gr.Radio(["CSV", "JSON Lines", "Parquet"], value="CSV", label="Format")
                exp_btn = gr.Button("Export current page", variant="primary")
            exp_note = gr.Markdown()
            exp_file = gr.File(label="Download")
            snippet = gr.Markdown()

    # ---------------- wiring ---------------- #
    connect_btn.click(connect, [dataset_id, token],
                      [config, split, info_md, readme_md, file_pick, meta]) \
        .then(file_table, meta, files_tbl) \
        .then(code_snippet, [meta, config, split, strategy], snippet)

    dataset_id.submit(connect, [dataset_id, token],
                      [config, split, info_md, readme_md, file_pick, meta]) \
        .then(file_table, meta, files_tbl)

    config.change(on_config_change, [config, meta], split)

    load_inputs = [meta, config, split, offset, limit, token, strategy, maxlen, state]
    load_btn.click(load_rows, load_inputs, [table, banner, state, cols, row_idx]) \
        .then(lambda s: gr.update(choices=[str(c) for c in _state_df(s).columns]), state, regex_col) \
        .then(media_choices, state, media_col) \
        .then(code_snippet, [meta, config, split, strategy], snippet)

    next_btn.click(lambda *a: page(1, *a), load_inputs,
                   [table, banner, state, cols, row_idx, offset])
    prev_btn.click(lambda *a: page(-1, *a), load_inputs,
                   [table, banner, state, cols, row_idx, offset])

    cols.change(pick_columns, [cols, state, maxlen], table)
    filter_btn.click(apply_filter, [query, regex_col, regex, state, maxlen], [ftable, filter_note])
    row_idx.change(inspect, [row_idx, state], [rec_json, rec_md, rec_img, rec_aud])
    stats_btn.click(profile, state, [schema_tbl, num_tbl, dist_md])
    media_btn.click(gallery, [state, media_col, n_media], [gal, media_note])
    raw_btn.click(read_raw, [meta, file_pick, raw_off, raw_lim, token, maxlen], [raw_tbl, raw_note])
    exp_btn.click(export, [state, fmt], [exp_file, exp_note])


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--share", action="store_true")
    p.add_argument("--port", type=int, default=7860)
    a = p.parse_args()
    kwargs = dict(server_name="0.0.0.0", server_port=a.port,
                  share=a.share, show_error=True)
    try:  # gradio >= 6 moved theme/css to launch()
        demo.queue(default_concurrency_limit=4).launch(
            theme=gr.themes.Soft(), css=CSS, **kwargs)
    except TypeError:
        demo.queue(default_concurrency_limit=4).launch(**kwargs)