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",
    "Avg Fair",
    "Gold69 v2 Fair CER",
    "Gold69 v2 Fair WER",
    "FLEURS WER",
    "FLEURS CER",
    "Decode (ms)",
    "Family",
    "Notes",
]

DATATYPES = [
    "markdown",
    "str",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "number",
    "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_wer(row: dict):
    gold = row.get("gold69_v2_fair_cer")
    fleurs = row.get("fleurs_cer")
    if gold is None or fleurs is None or pd.isna(gold) or pd.isna(fleurs):
        return None
    return (float(gold) + float(fleurs)) / 2


def model_link(row: dict) -> str:
    model = row.get("model", "")
    repo = row.get("repo", "")
    if isinstance(repo, str) and repo.startswith("http"):
        return f"[{model}]({repo})"
    return 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)

    records = []
    for row in merged:
        records.append(
            {
                "Model": model_link(row),
                "Params": num(row.get("params_b"), 3),
                "Avg Fair": pct(avg_wer(row)),
                "Gold69 v2 Fair CER": pct(row.get("gold69_v2_fair_cer")),
                "Gold69 v2 Fair WER": pct(row.get("gold69_v2_fair_wer")),
                "FLEURS WER": pct(row.get("fleurs_wer")),
                "FLEURS CER": pct(row.get("fleurs_cer")),
                "Decode (ms)": num(row.get("mean_decode_ms"), 1),
                "Family": row.get("family", ""),
                "Status": row.get("status", "queued"),
                "Notes": row.get("notes", ""),
                "_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("Avg Fair", ascending=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()
    completed = df[df["Status"].eq("complete")]
    best_gold = completed["Gold69 v2 Fair CER"].min() if not completed.empty else None
    best_fleurs = completed["FLEURS WER"].min() if not completed.empty else None
    return f"""
<div class="cards">
  <div><span>Completed models</span><strong>{len(completed)}</strong></div>
  <div><span>Best Gold69 v2 fair CER</span><strong>{best_gold:.2f}%</strong></div>
  <div><span>Best FLEURS WER</span><strong>{best_fleurs:.2f}%</strong></div>
  <div><span>Queued baselines</span><strong>{len(df) - len(completed)}</strong></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: 1220px !important;
  margin: 0 auto !important;
}
.hero h1 {
  font-size: clamp(34px, 5vw, 58px);
  line-height: 1.05;
  margin-bottom: 8px;
}
.hero p {
  color: #f0b7bd;
  font-size: 18px;
  line-height: 1.55;
}
.cards {
  display: grid;
  grid-template-columns: repeat(4, minmax(0, 1fr));
  gap: 10px;
  margin: 12px 0 18px;
}
.cards div {
  border: 1px solid #5b2430;
  background: linear-gradient(180deg, #35151c, #241014);
  border-radius: 8px;
  padding: 14px 16px;
}
.cards span {
  display: block;
  color: #f0b7bd;
  margin-bottom: 6px;
}
.cards strong {
  font-size: 28px;
}
#leaderboard table {
  table-layout: auto !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;
}
#leaderboard td {
  white-space: normal !important;
  overflow-wrap: anywhere !important;
}
#leaderboard td a {
  color: #1d4ed8 !important;
  font-weight: 700;
}
@media (max-width: 760px) {
  .cards { grid-template-columns: 1fr; }
  .hero h1 { font-size: 32px; }
}
"""


with gr.Blocks(title="Persian ASR Double Benchmark", css=CSS, theme=gr.themes.Soft(primary_hue="red", neutral_hue="rose")) as demo:
    gr.HTML(
        """
<section class="hero">
  <h1>Persian ASR Double Benchmark</h1>
  <p>
    A same-test leaderboard for Persian speech recognition models. Gold69 v2 is the corrected tougher
    real-world set; FLEURS is the cleaner public reference set. Gold69 v2 fair CER ignores punctuation,
    spaces, half-spaces, and diacritics so models are judged on the words they heard. Lower is better.
  </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",
        elem_id="leaderboard",
        max_height=760,
        show_search="none",
        show_copy_button=True,
        show_fullscreen_button=True,
        show_row_numbers=True,
        pinned_columns=1,
        max_chars=120,
        column_widths=[270, 100, 110, 116, 116, 116, 116, 132, 190, 420],
    )


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