"""Build Your Own SafeTulu: pick a removal threshold and sources, see what gets removed, export."""

from __future__ import annotations

import os
import tempfile
import time
from pathlib import Path
from urllib.parse import urlencode

import gradio as gr
import numpy as np
import plotly.graph_objects as go

import labs_theme as lt
import safetulu_data as sd

# ZeroGPU Spaces refuse to start unless at least one function is decorated with
# @spaces.GPU. This app needs no GPU, so register a no-op that is never called;
# everything else runs on the CPU side as usual. Elsewhere this is skipped.
try:
    import spaces

    @spaces.GPU(duration=5)
    def _zerogpu_placeholder():
        return None
except ImportError:
    pass

DF = sd.load_scores()
SOURCES = DF["source"].astype(str).value_counts().index.tolist()  # largest first
SOURCE_N = DF["source"].astype(str).value_counts().to_dict()
def short(name, n=40):
    return name if len(name) <= n else name[: n // 2 - 1] + "…" + name[-(n // 2 - 1):]


SOURCE_CHOICES = [(f"{short(s)}  ({SOURCE_N[s]:,})", s) for s in SOURCES]
SCORE_MIN = float(np.floor(DF[sd.SCORE_COL].min() * 1000) / 1000)
SCORE_MAX = float(np.ceil(DF[sd.SCORE_COL].max() * 1000) / 1000)
OFFICIAL_PCT = 100.0 * DF["official_removed"].mean()
SPACE_ID = os.environ.get("SPACE_ID")  # set automatically on HF Spaces
ON_SPACES = bool(SPACE_ID) or os.environ.get("SAFETULU_ENABLE_PUSH") == "1"

# Colors (validated categorical slots 1 and 2); text stays neutral.
KEPT, REMOVED = lt.KEPT, lt.REMOVED
INK, GRID = lt.INK, lt.GRID


# ---------------------------------------------------------------------------
# Figures
# ---------------------------------------------------------------------------
def _layout(fig, height, **kw):
    fig.update_layout(
        height=height,
        margin=dict(l=8, r=8, t=8, b=8),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(color=INK, size=12, family=f"{lt.FONT}, ui-sans-serif, sans-serif"),
        hoverlabel=dict(font_size=12, bgcolor=lt.NAVY_2, bordercolor=lt.ELECTRIC,
                        font=dict(color=lt.TEXT, family=f"{lt.FONT}, sans-serif")),
        legend=dict(traceorder="normal", orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        **kw,
    )
    fig.update_xaxes(gridcolor=GRID, zeroline=False, linecolor=GRID)
    fig.update_yaxes(gridcolor=GRID, zeroline=False, linecolor=GRID)
    return fig


def hist_figure(in_pool, removed, cutoff, log_y):
    edges, kept_c, rem_c = sd.histogram(DF, in_pool, removed)
    mids = (edges[:-1] + edges[1:]) / 2
    width = edges[1] - edges[0]
    ranges = [f"{a:.3f} – {b:.3f}" for a, b in zip(edges[:-1], edges[1:])]
    fig = go.Figure()
    for name, counts, color in (("Kept", kept_c, KEPT), ("Removed", rem_c, REMOVED)):
        fig.add_bar(
            x=mids, y=counts, width=width * 0.92, name=name, marker_color=color,
            customdata=ranges,
            hovertemplate=f"<b>{name}</b><br>score %{{customdata}}<br>%{{y:,}} examples<extra></extra>",
        )
    if cutoff is not None:
        fig.add_vline(x=cutoff, line=dict(color=lt.CUTOFF_LINE, width=1, dash="dot"))
        fig.add_annotation(x=cutoff, y=1, yref="paper", text=f"cutoff {cutoff:.3f}", showarrow=False,
                           xanchor="left", yanchor="top", xshift=4, font=dict(color=lt.TEXT, size=11))
    fig = _layout(fig, 340, barmode="stack", bargap=0)
    fig.update_layout(margin=dict(t=36))
    fig.update_xaxes(title_text="score (higher = closer to SafeRLHF harmful data)")
    fig.update_yaxes(title_text="examples", type="log" if log_y else "linear")
    return fig


def source_figure(table):
    t = table.sort_values("pct_removed", ascending=True)
    fig = go.Figure(go.Bar(
        y=[short(x) for x in t.source], x=t.pct_removed, orientation="h", marker_color=REMOVED,
        customdata=np.stack([t.removed, t.examples, t.mean_score], axis=1),
        hovertemplate="<b>%{y}</b><br>%{x:.1f}% removed<br>%{customdata[0]:,} of %{customdata[1]:,}"
                      "<br>mean score %{customdata[2]:.3f}<extra></extra>",
        text=[f"{v:.1f}%" for v in t.pct_removed], textposition="outside", cliponaxis=False,
        textfont=dict(color=INK, size=11),
    ))
    fig = _layout(fig, max(220, 26 * len(t) + 60), showlegend=False)
    fig.update_xaxes(title_text="% of source removed", ticksuffix="%", range=[0, max(1.0, float(t.pct_removed.max())) * 1.15])
    fig.update_yaxes(automargin=True)
    return fig


# ---------------------------------------------------------------------------
# Main update
# ---------------------------------------------------------------------------
def settings_from(mode_label, pct, cutoff, include, exempt) -> sd.Settings:
    mode = next(k for k, v in sd.MODES.items() if v == mode_label)
    inc = None if set(include or []) == set(SOURCES) else list(include or [])
    return sd.Settings(mode=mode, pct=float(pct), cutoff=float(cutoff), include=inc, exempt=list(exempt or []))


def stat_tiles(st, s):
    def tile(label, value, sub="", cls=""):
        return (f'<div class="tile {cls}"><div class="tile-label">{label}</div>'
                f'<div class="tile-value">{value}</div><div class="tile-sub">{sub}</div></div>')

    cut = "n/a" if st["score_cutoff"] is None else f"{st['score_cutoff']:.4f}"
    jac = st["jaccard_with_official"]
    agree = "n/a" if jac is None else f"{100 * jac:.0f}%"
    return '<div class="tiles">' + "".join([
        tile("Kept", f"{st['kept']:,}", f"of {st['pool']:,} in the selected sources"),
        tile("Removed", f"{st['removed']:,}", f"{st['removed_pct_of_pool']:.1f}% of the selection", "accent"),
        tile("Lowest removed score", cut, "examples at or above this are removed" if s.mode != "per_source"
             else "varies by source in this mode"),
        tile("Agreement with released SafeTulu", agree,
             f"Jaccard overlap of removed sets ({st['overlap_with_official']:,} shared)"),
    ]) + "</div>"


def update(mode_label, pct, cutoff, include, exempt, log_y):
    s = settings_from(mode_label, pct, cutoff, include, exempt)
    if not include:
        empty = '<div class="tiles"><div class="tile"><div class="tile-label">No sources selected</div></div></div>'
        return empty, go.Figure(), go.Figure(), None, share_link(s), sd.snippet(s)
    in_pool, removed = sd.select(DF, s)
    st = sd.summarize(DF, in_pool, removed)
    table = sd.per_source(DF, in_pool, removed)
    shown = table.rename(columns={"pct_removed": "% removed", "mean_score": "mean score"}).round(
        {"% removed": 1, "mean score": 4})
    return (stat_tiles(st, s), hist_figure(in_pool, removed, st["score_cutoff"] if s.mode != "per_source" else None, log_y),
            source_figure(table), shown, share_link(s), sd.snippet(s))


def share_link(s: sd.Settings) -> str:
    qs = urlencode(s.to_query())
    base = f"https://huggingface.co/spaces/{SPACE_ID}" if SPACE_ID else ""
    return f"{base}?{qs}"


def mode_visibility(mode_label):
    is_cut = mode_label == sd.MODES["cutoff"]
    return gr.update(visible=not is_cut), gr.update(visible=is_cut)


def on_load(request: gr.Request):
    q = dict(request.query_params) if request else {}
    s = sd.Settings.from_query(q, SOURCES)
    include = SOURCES if s.include is None else s.include
    is_cut = s.mode == "cutoff"
    controls_v = (sd.MODES[s.mode], s.pct, s.cutoff, include, s.exempt)
    return (controls_v[0], gr.update(value=s.pct, visible=not is_cut), gr.update(value=s.cutoff, visible=is_cut),
            include, s.exempt, gr.update(open=bool(s.exempt)), *update(*controls_v, False))


# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------
def export_ids(mode_label, pct, cutoff, include, exempt):
    s = settings_from(mode_label, pct, cutoff, include, exempt)
    in_pool, removed = sd.select(DF, s)
    out = Path(tempfile.mkdtemp(prefix="safetulu_ids_"))
    return [str(p) for p in sd.write_id_files(DF, in_pool, removed, s, out)]


def export_hub(mode_label, pct, cutoff, include, exempt, repo_name, private, with_scores, with_emb,
               oauth_token: gr.OAuthToken | None, profile: gr.OAuthProfile | None,
               progress=gr.Progress()):
    if oauth_token is None or profile is None:
        raise gr.Error("Sign in with Hugging Face first.")
    repo_name = (repo_name or "").strip().strip("/")
    if not repo_name:
        raise gr.Error("Enter a dataset name.")
    repo_id = repo_name if "/" in repo_name else f"{profile.username}/{repo_name}"
    s = settings_from(mode_label, pct, cutoff, include, exempt)
    in_pool, removed = sd.select(DF, s)
    st = sd.summarize(DF, in_pool, removed)
    keep_ids = DF[sd.ID_COL].to_numpy()[in_pool & ~removed]
    t0 = time.time()
    work = Path(tempfile.mkdtemp(prefix="safetulu_push_"))
    files = sd.materialize(keep_ids, work / "data", include_scores=with_scores,
                           include_embeddings=with_emb, progress=progress)
    progress(0.95, desc="Uploading to the Hub")
    url = sd.push_to_hub(files, sd.dataset_card(s, st), repo_id, oauth_token.token, private=private)
    return f"Pushed {st['kept']:,} examples to [{repo_id}]({url}) in {time.time() - t0:.0f}s."


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Build Your Own SafeTulu") as demo:
    gr.HTML(lt.hero_html(sd.REPO_ID, len(DF), OFFICIAL_PCT))
    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=300):
            mode = gr.Radio(list(sd.MODES.values()), value=sd.MODES["global"], label="How to remove")
            pct = gr.Slider(0, 50, value=20, step=0.5, label="Remove (%)",
                            info="Share of the filterable examples to remove")
            cutoff = gr.Slider(SCORE_MIN, SCORE_MAX, value=round((SCORE_MIN + SCORE_MAX) / 2, 4), step=0.0005,
                               label="Score cutoff", info="Remove examples with score ≥ this", visible=False)
            include = gr.CheckboxGroup(SOURCE_CHOICES, value=SOURCES, label="Sources to include in the output")
            with gr.Row():
                all_btn = gr.Button("Select all", size="sm")
                none_btn = gr.Button("Clear", size="sm")
            with gr.Accordion("Never filter these sources", open=False) as exempt_acc:
                exempt = gr.CheckboxGroup(SOURCE_CHOICES, value=[], show_label=False,
                                          info="Kept in full; the removal % applies to the rest")
            log_y = gr.Checkbox(False, label="Log scale for the histogram")
        with gr.Column(scale=2):
            tiles = gr.HTML()
            gr.Markdown("Score distribution", elem_classes="section-title")
            hist = gr.Plot(show_label=False)
            gr.Markdown("Share removed, by source", elem_classes="section-title")
            src_plot = gr.Plot(show_label=False)
            with gr.Accordion("Per-source table", open=False):
                table = gr.Dataframe(interactive=False, wrap=True)

    with gr.Tab("Share"):
        link = gr.Textbox(label="Link to this exact selection", buttons=["copy"], interactive=False)
    with gr.Tab("Download IDs"):
        gr.Markdown("Parquet files of kept and removed `document_id`s (with source and score), plus the settings as JSON.")
        ids_btn = gr.Button("Generate files", variant="primary")
        ids_files = gr.File(label="Files", file_count="multiple")
    with gr.Tab("Python script"):
        gr.Markdown("Reproduces this selection from the Hub without the Space and builds a `datasets.Dataset` of the kept examples.")
        code = gr.Code(language="python", label="Reproduce", max_lines=80, wrap_lines=True)
    with gr.Tab("Push to Hub"):
        gr.Markdown("Writes the kept examples (`document_id`, `content_hash`, `messages`) to a dataset repo in your account, with a dataset card recording these settings. Takes a few minutes.")
        if ON_SPACES:
            gr.LoginButton()
            repo_name = gr.Textbox(label="Dataset name", placeholder="my-safetulu-10pct")
            with gr.Row():
                private = gr.Checkbox(True, label="Private")
                with_scores = gr.Checkbox(False, label="Include score columns")
                with_emb = gr.Checkbox(False, label="Include gradient embeddings (adds ~7 GB)")
            push_btn = gr.Button("Push to Hub", variant="primary")
            push_status = gr.Markdown()
        else:
            gr.Markdown("*Sign-in only works on the hosted Space. Running locally, use the Python script tab and `ds.push_to_hub(...)`.*")

    controls = [mode, pct, cutoff, include, exempt]
    outputs = [tiles, hist, src_plot, table, link, code]
    for c in controls + [log_y]:
        c.input(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden")
    pct.release(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden")
    cutoff.release(update, controls + [log_y], outputs, trigger_mode="always_last", show_progress="hidden")
    mode.input(mode_visibility, mode, [pct, cutoff])
    all_btn.click(lambda: SOURCES, None, include).then(update, controls + [log_y], outputs, show_progress="hidden")
    none_btn.click(lambda: [], None, include).then(update, controls + [log_y], outputs, show_progress="hidden")
    ids_btn.click(export_ids, controls, ids_files)
    if ON_SPACES:
        push_btn.click(export_hub, controls + [repo_name, private, with_scores, with_emb], push_status)
    demo.load(on_load, None, controls + [exempt_acc] + outputs)

if __name__ == "__main__":
    demo.launch(css=lt.CSS, theme=lt.build_theme(), js=lt.THEME_JS)