import json
from pathlib import Path

import gradio as gr
import pandas as pd


ROOT = Path(__file__).parent
RESULTS_PATH = ROOT / "results.json"
MODELS_PATH = ROOT / "models.json"

COLUMNS = [
    "Model",
    "Params",
    "Triple Threat 🔥",
    "S³",
    "Ess Err 🔑",
    "VisualEars6669 Fair WER",
    "VisualEars6669 Fair CER",
    "FLEURS-fa Full WER",
    "FLEURS-fa Full CER",
    "Decode (ms)",
    "Family",
    "Status",
    "Notes",
]

DATATYPES = [
    "markdown", "str", "number", "number", "number", "number", "number",
    "number", "number", "str", "str", "str", "str",
]

STATUS_RANK = {"complete": 0, "running": 1, "queued": 2, "failed": 3}


def read_json(path: Path) -> list[dict]:
    if not path.exists():
        return []
    return json.loads(path.read_text())


def pct(value):
    return None if value is None or pd.isna(value) else round(float(value), 2)


def num(value, digits=3):
    return None if value is None or pd.isna(value) else round(float(value), digits)


def _avg(*xs):
    xs = [float(x) for x in xs if x is not None and not pd.isna(x)]
    return sum(xs) / len(xs) if xs else None


def s3_avg(row: dict):
    return _avg(row.get("visualears6669_s3"), row.get("fleurs_s3"))


def triple_threat(row: dict):
    """60% S³ + 20% WER + 20% CER, with each pillar split 50/50 by dataset.

    The score is defined only when all six dataset-level inputs exist, so a model
    is never flattered by a missing penalty term or an incomplete benchmark split.
    """
    s3 = s3_avg(row)
    p_wer = _avg(row.get("visualears6669_fair_wer"), row.get("fleurs_wer"))
    p_cer = _avg(row.get("visualears6669_fair_cer"), row.get("fleurs_cer"))
    required = (
        row.get("visualears6669_s3"), row.get("fleurs_s3"),
        row.get("visualears6669_fair_wer"), row.get("fleurs_wer"),
        row.get("visualears6669_fair_cer"), row.get("fleurs_cer"),
    )
    if any(value is None or pd.isna(value) for value in required):
        return None
    return 0.60 * s3 + 0.20 * p_wer + 0.20 * p_cer


def auto_note(row: dict) -> str:
    note = row.get("notes", "") or ""
    flags = []
    if s3_avg(row) is None and row.get("visualears6669_fair_wer") is not None:
        flags.append("S³ pending")
    elif row.get("fleurs_s3") is None and row.get("visualears6669_s3") is not None:
        flags.append("FLEURS S³ pending")
    if flags:
        tag = " · ".join(flags)
        note = f"{tag} — {note}" if note else tag
    return note


def medal_for_rank(rank: int | None) -> str:
    return {1: "🥇", 2: "🥈", 3: "🥉"}.get(rank, "")


def model_link(row: dict, medal: str = "") -> str:
    model = row.get("model", "")
    repo = row.get("repo", "")
    prefix = "🔒 " if row.get("closed_model") else ""
    medal_prefix = f"{medal} " if medal else ""
    if isinstance(repo, str) and repo.startswith("http"):
        return f"{medal_prefix}{prefix}[{model}]({repo})"
    return f"{medal_prefix}{prefix}{model}"


def combined_rows() -> pd.DataFrame:
    result_rows = {row.get("model"): row for row in read_json(RESULTS_PATH)}
    model_rows = read_json(MODELS_PATH)

    merged = []
    seen = set()
    for model_row in model_rows:
        model = model_row.get("model")
        merged.append({**model_row, **result_rows.get(model, {})})
        seen.add(model)
    for model, row in result_rows.items():
        if model not in seen:
            merged.append(row)

    tt_scores = {row.get("model"): triple_threat(row) for row in merged}
    ranked = sorted(
        ((m, s) for m, s in tt_scores.items() if s is not None and not pd.isna(s)),
        key=lambda item: (item[1], item[0] or ""),
    )
    medal_ranks = {m: rank for rank, (m, _s) in enumerate(ranked[:3], 1)}

    records = []
    for row in merged:
        records.append(
            {
                "Model": model_link(row, medal_for_rank(medal_ranks.get(row.get("model")))),
                "Params": num(row.get("params_b"), 3),
                "Triple Threat 🔥": pct(triple_threat(row)),
                "S³": pct(s3_avg(row)),
                "Ess Err 🔑": pct(row.get("ess_err")),
                "VisualEars6669 Fair WER": pct(row.get("visualears6669_fair_wer")),
                "VisualEars6669 Fair CER": pct(row.get("visualears6669_fair_cer")),
                "FLEURS-fa Full WER": pct(row.get("fleurs_wer")),
                "FLEURS-fa Full CER": pct(row.get("fleurs_cer")),
                "Decode (ms)": ("1×RT" if "Web Speech" in (row.get("model") or "") else ("" if (_d := num(row.get("mean_decode_ms"), 1)) is None else f"{_d}")),
                "Family": row.get("family", ""),
                "Status": row.get("status", "queued"),
                "Notes": auto_note(row),
                "_status_rank": STATUS_RANK.get(row.get("status", "queued"), 9),
                "_model_text": row.get("model", ""),
            }
        )

    return pd.DataFrame(records)


def display_frame() -> pd.DataFrame:
    df = combined_rows()
    df = df.sort_values(
        ["Triple Threat 🔥", "S³", "VisualEars6669 Fair WER"],
        ascending=[True, True, True],
        na_position="last",
    )
    display = df[COLUMNS].reset_index(drop=True).astype(object)
    return display.where(pd.notna(display), None)


def summary_cards() -> str:
    df = combined_rows()
    scored_tt = df[df["Triple Threat 🔥"].notna()]
    scored_s3 = df[df["S³"].notna()]

    def champ(frame, metric, icon):
        if frame.empty:
            return "—"
        r = frame.sort_values([metric, "_model_text"], ascending=[True, True]).iloc[0]
        return f"{icon} {r['_model_text']} · {float(r[metric]):.2f}"

    return f"""
<div class="cards">
  <div class="card hot">
    <span>🔥 Best Triple Threat</span>
    <strong>{champ(scored_tt, "Triple Threat 🔥", "🥇")}</strong>
    <em>60% S³ · 20% WER · 20% CER</em>
  </div>
  <div class="card sem">
    <span>🧠 Best S³ (semantic, DHH-weighted)</span>
    <strong>{champ(scored_s3, "S³", "🥇")}</strong>
    <em>lower = fewer meaning-critical errors</em>
  </div>
</div>
"""


CSS = """
:root {
  --body-background-fill: #12070a;
  --body-text-color: #fff7f7;
  --block-background-fill: #241014;
  --block-border-color: #5b2430;
  --button-primary-background-fill: #ef4444;
  --button-primary-background-fill-hover: #dc2626;
}
.gradio-container { max-width: 1260px !important; margin: 0 auto !important; }
.hero h1 { font-size: clamp(23px, 3.4vw, 40px); line-height: 1.1; white-space: nowrap; margin-bottom: 8px; }
.hero h1 .s3 { color:#fca5a5; }
.hero p { color: #f0b7bd; font-size: 17px; line-height: 1.58; }
.hero .pillars { margin: 10px 0 2px; display:flex; gap:8px; flex-wrap:wrap; }
.hero .pill { border:1px solid #7f1d1d; background:#2a0d12; color:#fecaca; border-radius:999px;
  padding:4px 12px; font-size:13px; font-weight:600; }
.hero .pill b { color:#fff; }
.cards { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 12px; margin: 14px 0 18px; }
.card { border: 1px solid #5b2430; background: linear-gradient(180deg, #35151c, #241014);
  border-radius: 10px; padding: 14px 16px; }
.card.hot { border-color:#ef4444; box-shadow: inset 0 0 0 1px rgba(239,68,68,.25); }
.card span { display: block; color: #f0b7bd; margin-bottom: 6px; font-size:14px; }
.card strong { font-size: 22px; display:block; }
.card em { color:#caa; font-size:12px; font-style:normal; }
#leaderboard table { table-layout: fixed !important; }
#leaderboard table, #leaderboard thead, #leaderboard tbody,
#leaderboard th, #leaderboard td, #leaderboard * { color: #2a090e !important; }
#leaderboard th, #leaderboard td { font-size: 15px !important; line-height: 1.32 !important; vertical-align: top !important; }
#leaderboard th { background: #ffe4e6 !important; color: #4c0519 !important; }
/* spotlight the Triple Threat + S³ + Essential-content columns (3rd–5th) */
#leaderboard th:nth-child(3), #leaderboard td:nth-child(3) { background:#fff1f2 !important; font-weight:700; }
#leaderboard th:nth-child(4), #leaderboard td:nth-child(4) { background:#fef2f2 !important; font-weight:700; }
#leaderboard th:nth-child(5), #leaderboard td:nth-child(5) { background:#fff5f5 !important; font-weight:700; }
#leaderboard td { white-space: normal !important; overflow-wrap: anywhere !important; }
#leaderboard td a { color: #1d4ed8 !important; font-weight: 700; }
#leaderboard [style*="max-height"],
#leaderboard .table-wrap,
#leaderboard .dataframe,
#leaderboard .wrap,
#leaderboard .scroll-hide,
#leaderboard .cell-wrap {
  max-height: none !important;
}
#leaderboard .table-wrap,
#leaderboard .dataframe,
#leaderboard .wrap,
#leaderboard .scroll-hide {
  overflow-y: visible !important;
}
/* user-adjustable column widths: drag the grip on the right edge of any header */
#leaderboard th { position: relative; }
#leaderboard th .col-resizer { position:absolute; top:0; right:0; width:9px; height:100%; cursor:col-resize; user-select:none; z-index:20; }
#leaderboard th .col-resizer::after { content:""; position:absolute; top:20%; right:3px; width:2px; height:60%; background:#f9a8b4; border-radius:2px; }
#leaderboard th .col-resizer:hover::after { background:#ef4444; }
@media (max-width: 760px) { .cards { grid-template-columns: 1fr; } .hero h1 { font-size: 24px; white-space: normal; } }
"""


COL_RESIZER_JS = """
() => {
  // Gradio renders two stacked tables (sticky-header layer + body layer) with no colgroup,
  // and the visible body does NOT inherit header-cell width. So pin the width on the header
  // cell AND every body cell of that column, across ALL #leaderboard tables. Persist widths
  // so a re-sort (which re-renders the rows) re-applies them.
  const widths = {};
  const applyWidth = (i, w) => {
    widths[i] = w;
    document.querySelectorAll('#leaderboard table').forEach((tb) => {
      const cells = tb.querySelectorAll(
        'thead th:nth-child(' + (i + 1) + '), tbody tr td:nth-child(' + (i + 1) + ')'
      );
      cells.forEach((c) => { c.style.width = w; c.style.minWidth = w; c.style.maxWidth = w; });
    });
  };
  const setup = () => {
    for (const i in widths) applyWidth(+i, widths[i]);
    document.querySelectorAll('#leaderboard table').forEach((table) => {
      table.querySelectorAll('thead th').forEach((th, i) => {
        if (i === 0) return;                      // skip the row-number column
        if (th.querySelector('.col-resizer')) return;
        const grip = document.createElement('div');
        grip.className = 'col-resizer';
        th.appendChild(grip);
        grip.addEventListener('mousedown', (e) => {
          e.preventDefault();
          e.stopPropagation();
          const startX = e.pageX;
          const startW = th.offsetWidth;
          const move = (ev) => applyWidth(i, Math.max(40, startW + ev.pageX - startX) + 'px');
          const up = () => {
            document.removeEventListener('mousemove', move);
            document.removeEventListener('mouseup', up);
          };
          document.addEventListener('mousemove', move);
          document.addEventListener('mouseup', up);
        });
      });
    });
  };
  setup();
  new MutationObserver(() => setup()).observe(document.body, { childList: true, subtree: true });
}
"""

with gr.Blocks(title="Persian ASR Triple Threat🔥 Benchmark", css=CSS, js=COL_RESIZER_JS,
               theme=gr.themes.Soft(primary_hue="red", neutral_hue="rose")) as demo:
    gr.HTML(
        """
<section class="hero">
  <h1>Persian ASR Triple Threat <span class="s3">🔥</span> Benchmark</h1>
  <p>
    A same-test leaderboard for Persian speech recognition. Two real-world sets — <b>VisualEars6669</b>
    (6,669-row / 10.49h noisy benchmark with clean, farfield &amp; obstructed conditions) and
    <b>FLEURS-fa Full</b> (4,341-row public Persian FLEURS) — now scored three ways:
  </p>
  <div class="pillars">
    <span class="pill"><b>WER</b> word error</span>
    <span class="pill"><b>CER</b> character error</span>
    <span class="pill"><b>S³</b> ShenavaSanj Score — semantic, DHH-weighted</span>
    <span class="pill"><b>🔑 Ess Err</b> essential-content error rate</span>
  </div>
  <p>
    <b>S³</b> is a Persian-first <b>Deaf/hard-of-hearing-focused semantic metric</b>: it weights every ASR
    error by how much that word matters for understanding, so dropping a keyword costs far more than a filler.
    <b>🔑 Ess Err</b> is the error rate on only the <b>most meaning-critical words</b> (importance ≥ 0.8) — the
    words a DHH reader can least afford to lose. <b>Triple Threat 🔥 = 60% S³ + 20% WER + 20% CER</b>.
    Each pillar is averaged 50/50 across VisualEars6669 and FLEURS-fa Full, so the two dataset splits remain equally weighted.
    Fair metrics ignore punctuation, spaces, half-spaces and diacritics, and normalize numbers
    (spelled-out ↔ digits) so a correctly-heard «صد» isn't penalized against «۱۰۰». Lower is better.
    🔒 marks closed paid APIs. S³ is undergoing validation with Persian DHH annotators.
  </p>
</section>
"""
    )
    gr.HTML(summary_cards)

    gr.Dataframe(
        value=display_frame,
        headers=COLUMNS,
        datatype=DATATYPES,
        interactive=False,
        wrap=True,
        line_breaks=True,
        label="Sortable leaderboard — sorted by Triple Threat",
        elem_id="leaderboard",
        max_height=1800,
        show_search="none",
        show_copy_button=True,
        show_fullscreen_button=True,
        show_row_numbers=True,
        pinned_columns=1,
        max_chars=120,
        column_widths=[240, 76, 100, 76, 92, 130, 130, 114, 114, 100, 160, 84, 280],
    )


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