import json
import os
from pathlib import Path, PurePosixPath

from huggingface_hub import HfApi, hf_hub_download
from shiny import App, reactive, render, ui

import coverage as cov
import inspect_coverage as icov

HARNESS_REPO = "MIMIR-AI-ROUTER/harness_evals"
INSPECT_REPO = "MIMIR-AI-ROUTER/inspect_evals"
HF_HARNESS = f"https://huggingface.co/datasets/{HARNESS_REPO}/blob/main"
HF_INSPECT = f"https://huggingface.co/datasets/{INSPECT_REPO}/blob/main"
HF_TOKEN = os.getenv("HF_TOKEN")
INSPECT_SUMMARY = Path(__file__).parent / "inspect_summary.json"

# ── CSS ───────────────────────────────────────────────────────────────────────

CSS = """
:root {
    --bg:   #1a1a2e; --bg2:  #16213e; --bg3: #0f3460;
    --acc:  #e94560; --link: #4a9eff;
    --text: #e0e0e0; --muted:#8888aa; --border:#2a2a4a;
    --ev:   #1e3a1e; --ev-t: #6fcf6f;
    --part: #2a2010; --part-t:#e0a020;
    --none-t:#555570; --skip-t:#666680;
}
* { box-sizing:border-box; margin:0; padding:0; }
body { background:var(--bg); color:var(--text); font-family:system-ui,sans-serif; }
.wrap { max-width:1500px; margin:0 auto; padding:20px; }
h1 { font-size:1.5em; border-bottom:2px solid var(--acc); padding-bottom:8px; margin-bottom:6px; }
.sub { color:var(--muted); font-size:.82em; margin-bottom:16px; }
.sub a { color:var(--link); text-decoration:none; }

/* top tabs */
.top-tabs { display:flex; gap:4px; border-bottom:2px solid var(--border); margin-bottom:18px; }
.top-tab { background:none; border:none; color:var(--muted); padding:8px 18px;
           cursor:pointer; font-size:.9em; border-bottom:2px solid transparent;
           margin-bottom:-2px; }
.top-tab:hover { color:var(--text); }
.top-tab.active { color:var(--text); border-bottom-color:var(--acc); }
.top-pane { display:none; }
.top-pane.active { display:block; }

/* controls */
.controls { display:flex; align-items:center; gap:12px; margin-bottom:14px; flex-wrap:wrap; }
.btn { background:var(--link); color:#fff; border:none; padding:6px 13px;
       border-radius:4px; cursor:pointer; font-size:.82em; }
.btn:hover { opacity:.85; }
.filter-bar { display:flex; align-items:center; gap:10px; flex-wrap:wrap;
              margin-bottom:10px; }
.filter-search { background:var(--bg2); color:var(--text); border:1px solid var(--border);
                 border-radius:4px; padding:5px 10px; font-size:.82em; width:200px; }
.filter-label { color:var(--muted); font-size:.8em; }
.filter-check { color:var(--muted); font-size:.8em; cursor:pointer;
                display:flex; align-items:center; gap:4px; }

/* stat chips */
.chips { display:flex; gap:10px; margin-bottom:14px; flex-wrap:wrap; }
.chip { background:var(--bg2); border-radius:5px; padding:8px 14px;
        font-size:.82em; border-left:3px solid var(--border); }
.chip .n { font-size:1.35em; font-weight:700; display:block; }
.chip-ev   { border-left-color:#4caf50; }
.chip-part { border-left-color:#ff9800; }
.chip-skip { border-left-color:var(--muted); }
.chip-none { border-left-color:var(--acc); }

/* matrix table */
.matrix-wrap { overflow-x:auto; max-height:80vh; overflow-y:auto; }
.matrix-table { border-collapse:collapse; width:100%; }
.matrix-table th { background:var(--bg2); color:var(--muted); font-size:.72em;
                   text-transform:uppercase; letter-spacing:.06em;
                   padding:8px 12px; text-align:left; border-bottom:2px solid var(--border);
                   position:sticky; top:0; white-space:nowrap; }
.matrix-table th.c-name-h { position:sticky; top:0; left:0; z-index:3; min-width:200px; }
.matrix-table th.c-val-h  { min-width:120px; text-align:center; }
.matrix-table td { padding:7px 10px; border-bottom:1px solid var(--border);
                   font-size:.84em; }
.matrix-table td.c-name { position:sticky; left:0; background:var(--bg); z-index:1;
                           min-width:200px; font-size:.83em; }
.matrix-table tbody tr:hover td { background:var(--bg2); }
.matrix-table tbody tr:hover td.c-name { background:var(--bg2); }

/* cells */
.c-ev   { background:var(--ev); text-align:center; }
.c-part { background:var(--part); text-align:center; }
.c-none { text-align:center; color:var(--none-t); }
.c-skip { text-align:center; color:var(--skip-t); font-size:.78em; font-style:italic; }
.c-na   { text-align:center; color:var(--muted); }
.val    { color:var(--ev-t); font-family:monospace; font-weight:600; font-size:.95em; }
.val-part { color:var(--part-t); font-family:monospace; font-size:.9em; }
.metric-lbl { color:var(--muted); font-size:.72em; display:block; }
.n-lbl  { color:var(--muted); font-size:.72em; display:block; }
.a { color:inherit; text-decoration:none; }
.a:hover { text-decoration:underline; }
.badge-skip { font-size:.72em; color:var(--skip-t); }
.err { color:var(--acc); background:rgba(233,69,96,.1); padding:10px 14px;
       border-radius:4px; border-left:3px solid var(--acc); margin-bottom:12px; }
.loading { color:var(--muted); padding:30px; text-align:center; font-style:italic; }
"""

_JS = """
function switchTop(name) {
    ['harness','inspect'].forEach(n => {
        document.getElementById('tab-' + n).classList.toggle('active', n === name);
        document.getElementById('pane-' + n).classList.toggle('active', n === name);
    });
}
function filterTable(tid) {
    const q = (document.getElementById(tid+'-q') || {value:''}).value.toLowerCase();
    const checks = {
        evaluated: (document.getElementById(tid+'-ev')   || {checked:true}).checked,
        partial:   (document.getElementById(tid+'-part') || {checked:true}).checked,
        not_started:(document.getElementById(tid+'-none')|| {checked:false}).checked,
        skipped:   (document.getElementById(tid+'-skip') || {checked:false}).checked,
    };
    document.querySelectorAll('#' + tid + ' tbody tr').forEach(r => {
        const ok = (!q || r.dataset.name.includes(q)) && (checks[r.dataset.status] !== false);
        r.style.display = ok ? '' : 'none';
    });
}
"""

# ── helpers ────────────────────────────────────────────────────────────────────

def _short_model(m: str) -> str:
    parts = m.split("/")
    base = parts[-1]
    if len(parts) >= 3:
        return f"{base} ({parts[0]})"
    return base


def _cell_ev(entry: dict, blob_url: str, partial: bool = False) -> ui.Tag:
    val = entry["value"]
    metric = (entry.get("metric") or "").split("/")[-1]
    n = entry.get("n", 0)
    href = f"{blob_url}/{entry['filename']}"
    cls = "c-part" if partial else "c-ev"
    vcls = "val-part" if partial else "val"
    return ui.tags.td(
        {"class": cls},
        ui.tags.a(
            {"class": "a", "href": href, "target": "_blank"},
            ui.span({"class": vcls}, f"{val:.3f}"),
        ),
        ui.span({"class": "metric-lbl"}, metric),
    )


def _cell_na(entry: dict, blob_url: str) -> ui.Tag:
    href = f"{blob_url}/{entry['filename']}"
    return ui.tags.td(
        {"class": "c-na"},
        ui.tags.a({"class": "a", "href": href, "target": "_blank"}, "n/a"),
    )


def _cell_none() -> ui.Tag:
    return ui.tags.td({"class": "c-none"}, "—")


def _cell_skip(reason: str) -> ui.Tag:
    return ui.tags.td({"class": "c-skip"}, reason or "skip")


def _render_matrix(
    matrix: dict,
    all_rows: list[str],
    models: list[str],
    cov_data: dict,
    blob_url: str,
    table_id: str,
    show_partial: bool = True,
) -> ui.Tag:
    """Render a complete matrix table with filter bar and stat chips."""

    stats = {"evaluated": 0, "partial": 0, "not_started": 0, "skipped": 0}
    for info in cov_data.values():
        s = info["status"]
        if s in stats:
            stats[s] += 1

    chips = ui.div(
        {"class": "chips"},
        ui.div({"class": "chip chip-ev"},   ui.span({"class": "n"}, str(stats["evaluated"])),   "Evaluated"),
        *([] if not show_partial else [
            ui.div({"class": "chip chip-part"}, ui.span({"class": "n"}, str(stats["partial"])),  "Partial"),
        ]),
        ui.div({"class": "chip chip-none"}, ui.span({"class": "n"}, str(stats["not_started"])), "Not started"),
        ui.div({"class": "chip chip-skip"}, ui.span({"class": "n"}, str(stats["skipped"])),     "Skipped"),
    )

    part_check = [] if not show_partial else [
        ui.tags.label(
            {"class": "filter-check"},
            ui.tags.input({"type": "checkbox", "checked": "", "id": table_id + "-part",
                           "onchange": f"filterTable('{table_id}')"}),
            " Partial",
        )
    ]

    filter_bar = ui.div(
        {"class": "filter-bar"},
        ui.tags.input({"type": "text", "class": "filter-search", "placeholder": "Filter by name…",
                       "id": table_id + "-q", "oninput": f"filterTable('{table_id}')"}),
        ui.span({"class": "filter-label"}, "Show:"),
        ui.tags.label({"class": "filter-check"},
            ui.tags.input({"type": "checkbox", "checked": "", "id": table_id + "-ev",
                           "onchange": f"filterTable('{table_id}')"}), " Evaluated"),
        *part_check,
        ui.tags.label({"class": "filter-check"},
            ui.tags.input({"type": "checkbox", "id": table_id + "-none",
                           "onchange": f"filterTable('{table_id}')"}), " Not started"),
        ui.tags.label({"class": "filter-check"},
            ui.tags.input({"type": "checkbox", "id": table_id + "-skip",
                           "onchange": f"filterTable('{table_id}')"}), " Skipped"),
    )

    header = ui.tags.thead(
        ui.tags.tr(
            ui.tags.th("Eval", {"class": "c-name-h"}),
            *[ui.tags.th(_short_model(m), {"class": "c-val-h"}) for m in models],
        )
    )

    rows = []
    for row_name in all_rows:
        info = cov_data.get(row_name, {"status": "not_started", "skip_reason": None})
        status = info["status"]
        skip_reason = info.get("skip_reason")

        # n_samples: take from first available model entry
        n = 0
        for m in models:
            e = matrix.get(row_name, {}).get(m)
            if e and e.get("n"):
                n = e["n"]
                break

        name_cell = ui.tags.td(
            {"class": "c-name"},
            row_name,
            *([] if not n else [ui.tags.br(), ui.span({"class": "n-lbl"}, f"n={n:,}")]),
        )

        cells: list[ui.Tag] = [name_cell]
        for model in models:
            entry = matrix.get(row_name, {}).get(model)
            if skip_reason is not None:
                cells.append(_cell_skip(skip_reason))
            elif entry is None:
                cells.append(_cell_none())
            elif entry.get("value") is None:
                cells.append(_cell_na(entry, blob_url))
            else:
                cells.append(_cell_ev(entry, blob_url, partial=(status == "partial")))

        rows.append(ui.tags.tr({"data-name": row_name, "data-status": status}, *cells))

    return ui.div(
        chips,
        filter_bar,
        ui.div(
            {"class": "matrix-wrap"},
            ui.tags.table(
                {"class": "matrix-table", "id": table_id},
                header,
                ui.tags.tbody(*rows),
            ),
        ),
    )


# ── harness data ──────────────────────────────────────────────────────────────

_hcache: dict = {"runs": None, "err": None}


def _parse_harness_task_metrics(td: dict) -> tuple[str | None, float | None]:
    for k, v in td.items():
        if k in ("name", "alias", "sample_len") or not isinstance(v, (int, float)):
            continue
        if "," in k:
            mname = k.split(",")[0]
            if not mname.endswith("_stderr"):
                return mname, v
    return None, None


def load_runs(force: bool = False) -> tuple[list, str | None]:
    if not force and _hcache["runs"] is not None:
        return _hcache["runs"], _hcache["err"]

    api = HfApi(token=HF_TOKEN)
    try:
        all_files = sorted(api.list_repo_files(HARNESS_REPO, repo_type="dataset"))
    except Exception as e:
        _hcache.update(runs=[], err=str(e))
        return [], str(e)

    result_files = [
        f for f in all_files
        if PurePosixPath(f).name.startswith("results_") and f.endswith(".json")
    ]
    sample_files = {
        f for f in all_files
        if PurePosixPath(f).name.startswith("samples_") and f.endswith(".jsonl")
    }

    runs = []
    for filename in sorted(result_files, reverse=True):
        try:
            local = hf_hub_download(HARNESS_REPO, filename, repo_type="dataset", token=HF_TOKEN)
            with open(local) as f:
                data = json.load(f)
        except Exception:
            continue

        p = PurePosixPath(filename)
        folder = p.parent.as_posix()
        ts = p.name.removeprefix("results_").removesuffix(".json")
        model = data.get("model_name", "unknown")

        tasks = []
        for task_name, td in data.get("results", {}).items():
            if not isinstance(td, dict):
                continue
            mname, mval = _parse_harness_task_metrics(td)
            samples_f = f"{folder}/samples_{task_name}_{ts}.jsonl"
            tasks.append({
                "name": task_name,
                "n": td.get("sample_len", 0),
                "metric": mname,
                "value": mval,
                "samples_file": samples_f if samples_f in sample_files else None,
            })

        runs.append({"model": model, "filename": filename, "tasks": tasks})

    _hcache.update(runs=runs, err=None)
    return runs, None


def load_harness_matrix(force: bool = False) -> tuple[dict, list[str], set[str], str | None]:
    """Returns (folder_matrix, sorted_models, evaluated_task_names, error)."""
    runs, err = load_runs(force)
    if err:
        return {}, [], set(), err

    folder_set = set(cov.ALL_FOLDERS)
    evaluated_tasks: set[str] = set()

    # folder -> model -> list of task entries
    raw: dict[str, dict[str, list]] = {}
    for run in runs:
        model = run["model"]
        for task in run["tasks"]:
            if task["value"] is None:
                continue
            evaluated_tasks.add(task["name"])
            folder = cov._folder_for_task(task["name"], folder_set)
            if folder is None:
                continue
            raw.setdefault(folder, {}).setdefault(model, []).append({
                "task": task["name"],
                "value": task["value"],
                "metric": task["metric"],
                "n": task["n"],
                "filename": run["filename"],
                "is_group": task["name"] == folder,
            })

    matrix: dict[str, dict] = {f: {} for f in cov.ALL_FOLDERS}
    all_models: set[str] = set()

    for folder, by_model in raw.items():
        for model, entries in by_model.items():
            all_models.add(model)
            group = [e for e in entries if e["is_group"]]
            best = group[-1] if group else entries[-1]
            is_partial = not bool(group)
            matrix[folder][model] = {
                "value": best["value"],
                "metric": best["metric"],
                "n": best["n"],
                "filename": best["filename"],
                "partial": is_partial,
            }

    return matrix, sorted(all_models), evaluated_tasks, None


# ── inspect data ──────────────────────────────────────────────────────────────

_icache: dict = {"matrix": None, "models": None, "tasks": None}


def _load_inspect_summary() -> dict[str, dict]:
    """Load inspect_summary.json as a filename→metrics lookup."""
    if not INSPECT_SUMMARY.exists():
        return {}
    try:
        data = json.loads(INSPECT_SUMMARY.read_text())
        return {
            e["filename"]: {
                "task":   e["task"],
                "model":  e["model"],
                "value":  e.get("primary_value"),
                "metric": e.get("primary_metric"),
                "n":      e.get("n_samples", 0),
            }
            for e in data.get("entries", [])
        }
    except Exception:
        return {}


def load_inspect_matrix() -> tuple[dict, list[str], set[str], str | None]:
    """Returns (task_matrix, sorted_models, evaluated_task_names, error).

    Lists .eval files live from HF every call (when cache is empty).
    Uses inspect_summary.json as a metrics cache for known files;
    falls back to header_only for files not yet in the summary.
    """
    if _icache["matrix"] is not None:
        return _icache["matrix"], _icache["models"], _icache["tasks"], None

    from inspect_ai.log import read_eval_log

    # Metrics cache: filename -> {task, model, value, metric, n}
    summary = _load_inspect_summary()

    api = HfApi(token=HF_TOKEN)
    try:
        all_files = sorted(api.list_repo_files(INSPECT_REPO, repo_type="dataset"))
    except Exception as e:
        return {}, [], set(), str(e)

    eval_files = [f for f in all_files if f.endswith(".eval")]

    matrix: dict[str, dict] = {}
    all_models: set[str] = set()
    evaluated: set[str] = set()

    for fname in eval_files:
        if fname in summary:
            e = summary[fname]
            task, model = e["task"], e["model"]
            entry = {"value": e["value"], "metric": e["metric"],
                     "n": e["n"], "filename": fname}
        else:
            # New file not yet in summary: read header only for task/model
            try:
                local = hf_hub_download(INSPECT_REPO, fname,
                                        repo_type="dataset", token=HF_TOKEN)
                log = read_eval_log(local, header_only=True)
                task  = log.eval.task.split("/")[-1]
                model = log.eval.model
                n     = log.eval.dataset.samples if log.eval.dataset else 0
            except Exception:
                continue
            entry = {"value": None, "metric": None, "n": n, "filename": fname}

        all_models.add(model)
        evaluated.add(task)
        # Latest file (sorted chronologically by filename) wins
        matrix.setdefault(task, {})[model] = entry

    models = sorted(all_models)
    _icache.update(matrix=matrix, models=models, tasks=evaluated)
    return matrix, models, evaluated, None


# ── UI ────────────────────────────────────────────────────────────────────────

app_ui = ui.page_fluid(
    ui.tags.head(ui.tags.style(CSS)),
    ui.div(
        {"class": "wrap"},
        ui.h1("LM Eval Results"),
        ui.div(
            {"class": "sub"},
            "Harness: ",
            ui.tags.a(HARNESS_REPO,
                      href=f"https://huggingface.co/datasets/{HARNESS_REPO}",
                      target="_blank"),
            "  ·  Inspect: ",
            ui.tags.a(INSPECT_REPO,
                      href=f"https://huggingface.co/datasets/{INSPECT_REPO}",
                      target="_blank"),
        ),
        ui.div(
            {"class": "top-tabs"},
            ui.tags.button("Harness", {"class": "top-tab active", "id": "tab-harness",
                                       "onclick": "switchTop('harness')"}),
            ui.tags.button("Inspect", {"class": "top-tab",         "id": "tab-inspect",
                                       "onclick": "switchTop('inspect')"}),
        ),
        ui.div(
            {"class": "top-pane active", "id": "pane-harness"},
            ui.div(
                {"class": "controls"},
                ui.input_action_button("refresh_h", "↺ Refresh", class_="btn"),
            ),
            ui.output_ui("harness_matrix"),
        ),
        ui.div(
            {"class": "top-pane", "id": "pane-inspect"},
            ui.div(
                {"class": "controls"},
                ui.input_action_button("refresh_i", "↺ Refresh", class_="btn"),
            ),
            ui.output_ui("inspect_matrix"),
        ),
        ui.tags.script(_JS),
    ),
)


# ── server ────────────────────────────────────────────────────────────────────

def server(input, output, session):

    @reactive.effect
    @reactive.event(input.refresh_h)
    def _ref_h():
        _hcache.update(runs=None, err=None)

    @reactive.effect
    @reactive.event(input.refresh_i)
    def _ref_i():
        _icache.update(matrix=None, models=None, tasks=None)

    @output
    @render.ui
    def harness_matrix():
        input.refresh_h()
        matrix, models, ev_tasks, err = load_harness_matrix()
        if err:
            return ui.div({"class": "err"}, f"Harness load error: {err}")
        if not models:
            return ui.div({"class": "loading"}, "No harness eval results found.")

        # Override coverage status to include partial from matrix
        coverage = cov.compute_coverage(ev_tasks)
        for folder in cov.ALL_FOLDERS:
            if coverage[folder]["status"] == "evaluated" and matrix.get(folder):
                for m_data in matrix[folder].values():
                    if m_data.get("partial"):
                        coverage[folder]["status"] = "partial"
                        break

        return _render_matrix(
            matrix, cov.ALL_FOLDERS, models, coverage, HF_HARNESS, "h-mat",
            show_partial=True,
        )

    @output
    @render.ui
    def inspect_matrix():
        input.refresh_i()
        matrix, models, ev_tasks, err = load_inspect_matrix()
        if err:
            return ui.div({"class": "err"}, f"Inspect load error: {err}")

        coverage = icov.compute_coverage(ev_tasks)
        # Rows = all inspect modules + any evaluated tasks not mapped to a module
        unmapped = sorted(ev_tasks - {
            t for m_tasks in coverage.values() for t in m_tasks["matched_tasks"]
        })
        all_rows = icov.ALL_MODULES + unmapped

        return _render_matrix(
            matrix, all_rows, models, coverage, HF_INSPECT, "i-mat",
            show_partial=False,
        )


app = App(app_ui, server)