"""Drupal AI Eval Dashboard — POC.

Reads EvalEval instance-level records from a HuggingFace dataset and
renders a filterable comparison dashboard across models and skills.
"""
import json
import shutil
import tempfile
from pathlib import Path

import gradio as gr
import pandas as pd
from huggingface_hub import snapshot_download

DATASET_REPO = "webchick/eval-results-poc"

CSS = """
.dashboard-header {
    background: linear-gradient(135deg, #004F8C 0%, #0678BE 100%);
    padding: 24px 32px;
    border-radius: 12px;
    color: white;
    margin-bottom: 16px;
}
.dashboard-header h1 {
    margin: 0;
    font-size: 1.8em;
    font-weight: 700;
    color: white;
}
.dashboard-header p {
    margin: 6px 0 0;
    opacity: 0.9;
    font-size: 0.95em;
    color: white;
}
.dashboard-header a { color: white !important; text-decoration: underline; }
.kpi-grid {
    display: grid;
    grid-template-columns: repeat(4, 1fr);
    gap: 12px;
    margin-bottom: 16px;
}
.kpi-card {
    background: white;
    border: 1px solid #e2e8f0;
    border-radius: 10px;
    padding: 16px;
    text-align: center;
    box-shadow: 0 1px 4px rgba(0,0,0,0.07);
}
.kpi-value {
    font-size: 2em;
    font-weight: 700;
    color: #0678BE;
    line-height: 1;
}
.kpi-label {
    font-size: 0.72em;
    color: #64748b;
    text-transform: uppercase;
    letter-spacing: 0.07em;
    margin-top: 6px;
}
"""

HEADER_HTML = """
<div class="dashboard-header">
    <h1>Drupal AI Eval Dashboard</h1>
    <p>How well do AI models handle Drupal-specific coding tasks?
    Results from <a href="https://www.drupal.org/project/ai_best_practices">ai_best_practices</a>,
    stored in <a href="https://github.com/evaleval/every_eval_ever">EvalEval envelope format</a>.</p>
</div>
"""


def load_results() -> pd.DataFrame:
    """Download the dataset fresh and flatten instance records into a DataFrame."""
    tmp = tempfile.mkdtemp()
    try:
        local_dir = snapshot_download(
            repo_id=DATASET_REPO,
            repo_type="dataset",
            local_dir=tmp,
        )
        rows = []
        for jsonl_file in sorted(Path(local_dir).rglob("*_samples.jsonl")):
            with open(jsonl_file) as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        rec = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    if not rec.get("schema_version", "").startswith("instance_level_eval"):
                        continue

                    meta  = rec.get("metadata") or {}
                    perf  = rec.get("performance") or {}
                    usage = rec.get("token_usage") or {}
                    ev    = rec.get("evaluation") or {}
                    inp   = rec.get("input") or {}
                    reference = inp.get("reference") or []
                    task = reference[0] if reference and reference[0] else inp.get("raw", "")[:100]

                    rows.append({
                        "Namespace":     meta.get("eval_namespace", "?"),
                        "Skill":         meta.get("skill_name", "?"),
                        "Case":          meta.get("case_id", "?"),
                        "Model":         rec.get("model_id", "?"),
                        "Passed":        bool(ev.get("is_correct")),
                        "Task":          task,
                        "Latency (ms)":  round(perf.get("latency_ms") or 0, 0),
                        "Input tokens":  usage.get("input_tokens", 0),
                        "Output tokens": usage.get("output_tokens", 0),
                        "Cost (USD)":    round(float(meta.get("cost_usd") or 0), 6),
                        "Run timestamp": meta.get("timestamp_run", ""),
                    })
        df = pd.DataFrame(rows)
        if not df.empty:
            df["Passed"] = df["Passed"].astype(bool)
        return df
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def make_kpi_html(df: pd.DataFrame) -> str:
    if df.empty:
        return ""
    total     = len(df)
    models    = df["Model"].nunique()
    skills    = df["Skill"].nunique()
    pass_rate = df["Passed"].mean() * 100

    def card(value, label):
        return (
            f'<div style="background:white;border:1px solid #e2e8f0;border-radius:10px;'
            f'padding:16px;text-align:center;box-shadow:0 1px 4px rgba(0,0,0,0.07)">'
            f'<div style="font-size:2em;font-weight:700;color:#0678BE;line-height:1">{value}</div>'
            f'<div style="font-size:0.72em;color:#475569;text-transform:uppercase;'
            f'letter-spacing:0.07em;margin-top:6px">{label}</div>'
            f'</div>'
        )

    return (
        f'<div style="display:grid;grid-template-columns:repeat(4,1fr);gap:12px;margin-bottom:16px">'
        f'{card(total, "Total runs")}'
        f'{card(models, "Models")}'
        f'{card(skills, "Skills")}'
        f'{card(f"{pass_rate:.0f}%", "Pass rate")}'
        f'</div>'
    )


def make_agg(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame(columns=["Model", "Pass rate (%)", "Cases run"])
    agg = (
        df.assign(_p=df["Passed"].astype(int))
        .groupby("Model")["_p"]
        .agg(["mean", "count"])
        .reset_index()
        .rename(columns={"mean": "Pass rate (%)", "count": "Cases run"})
    )
    agg["Pass rate (%)"] = (agg["Pass rate (%)"] * 100).round(1)
    return agg


# Numeric columns where lower = better (green) and higher = worse (red)
_LOWER_IS_BETTER = ["Latency (ms)", "Input tokens", "Output tokens", "Cost (USD)"]

_TABLE_CSS = """
<style>
.eval-table { width:100%; border-collapse:collapse; font-size:0.85em; }
.eval-table th {
    background:#0678BE; color:white; padding:8px 10px;
    text-align:left; white-space:nowrap; position:sticky; top:0;
}
.eval-table td { padding:6px 10px; border-bottom:1px solid #e2e8f0; color:#1e293b; }
.eval-table tr:hover td { background:#f0f9ff !important; color:#1e293b !important; }
.eval-table td.pass-cell { text-align:center; font-size:1.1em; }
</style>
"""


def _heat_rgb(ratio: float) -> str:
    """Map 0→green, 1→red as pastel background colours."""
    r = int(198 + (255 - 198) * ratio)
    g = int(239 - (239 - 199) * ratio)
    b = int(206 + (206 - 206) * ratio)
    return f"rgb({r},{g},{b})"


def df_to_html_table(df: pd.DataFrame) -> str:
    """Render DataFrame as an HTML table with emoji pass/fail and heat-map colours."""
    if df.empty:
        return "<p style='color:#64748b;padding:16px'>No results found.</p>"

    # Pre-compute min/max for numeric gradient columns
    ranges: dict[str, tuple[float, float]] = {}
    for col in _LOWER_IS_BETTER:
        if col in df.columns:
            mn, mx = float(df[col].min()), float(df[col].max())
            if mx > mn:
                ranges[col] = (mn, mx)

    cols = list(df.columns)

    # Header row
    ths = "".join(f"<th>{c}</th>" for c in cols)

    # Data rows
    rows_html = ""
    for _, row in df.iterrows():
        cells = ""
        for col in cols:
            val = row[col]
            if col == "Passed":
                icon = "✅" if val else "❌"
                cells += f'<td class="pass-cell">{icon}</td>'
            elif col in ranges:
                mn, mx = ranges[col]
                ratio = (float(val) - mn) / (mx - mn)
                bg = _heat_rgb(ratio)
                cells += (
                    f'<td style="background:{bg};color:#1e293b;white-space:nowrap">'
                    f'{val}</td>'
                )
            else:
                safe = str(val).replace("<", "&lt;").replace(">", "&gt;")
                cells += f"<td>{safe}</td>"
        rows_html += f"<tr>{cells}</tr>"

    return (
        _TABLE_CSS
        + f'<div style="overflow-x:auto"><table class="eval-table">'
        + f"<thead><tr>{ths}</tr></thead>"
        + f"<tbody>{rows_html}</tbody>"
        + "</table></div>"
    )


def apply_filters(df: pd.DataFrame, ns: str, skill: str, result: str) -> pd.DataFrame:
    filtered = df.copy()
    if ns != "(all)":
        filtered = filtered[filtered["Namespace"] == ns]
    if skill != "(all)":
        filtered = filtered[filtered["Skill"] == skill]
    if result == "Passed only":
        filtered = filtered[filtered["Passed"]]
    elif result == "Failed only":
        filtered = filtered[~filtered["Passed"]]
    return filtered


def build_ui():
    with gr.Blocks(
        title="Drupal AI Eval Dashboard",
        css=CSS,
        theme=gr.themes.Soft(),
    ) as demo:

        gr.HTML(HEADER_HTML)
        kpi_html = gr.HTML()

        with gr.Row():
            namespace_filter = gr.Dropdown(
                choices=["(all)"], value="(all)", label="Namespace", scale=2,
            )
            skill_filter = gr.Dropdown(
                choices=["(all)"], value="(all)", label="Skill", scale=2,
            )
            result_filter = gr.Radio(
                choices=["All", "Passed only", "Failed only"],
                value="All",
                label="Result",
                scale=1,
            )

        results_table = gr.HTML()

        gr.Markdown("### Pass rate by model")
        pass_rate_chart = gr.BarPlot(
            x="Model",
            y="Pass rate (%)",
            tooltip=["Model", "Pass rate (%)", "Cases run"],
            y_lim=[0, 100],
            color="Model",
        )

        data_store = gr.State(pd.DataFrame())

        def on_load():
            df = load_results()
            if df.empty:
                return (
                    gr.Dropdown(choices=["(all)"], value="(all)"),
                    gr.Dropdown(choices=["(all)"], value="(all)"),
                    df, "", df, pd.DataFrame(),
                )
            namespaces = sorted(df["Namespace"].unique().tolist())
            skills     = sorted(df["Skill"].unique().tolist())
            return (
                gr.Dropdown(choices=["(all)"] + namespaces, value="(all)"),
                gr.Dropdown(choices=["(all)"] + skills,      value="(all)"),
                df,
                make_kpi_html(df),
                df_to_html_table(df),
                make_agg(df),
            )

        def filter_data(ns, skill, result, df):
            filtered = apply_filters(df, ns, skill, result)
            if filtered.empty:
                return "", "<p style='color:#64748b;padding:16px'>No results found.</p>", pd.DataFrame()
            return make_kpi_html(filtered), df_to_html_table(filtered), make_agg(filtered)

        demo.load(
            fn=on_load,
            outputs=[
                namespace_filter, skill_filter,
                data_store, kpi_html,
                results_table, pass_rate_chart,
            ],
        )

        for widget in [namespace_filter, skill_filter, result_filter]:
            widget.change(
                fn=filter_data,
                inputs=[namespace_filter, skill_filter, result_filter, data_store],
                outputs=[kpi_html, results_table, pass_rate_chart],
            )

    return demo


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