"""Self-contained HTML report (all figures embedded as base64 PNG)."""

from __future__ import annotations

import base64
import html
import io
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from .calibration import reliability_bins  # noqa: E402

# Small fixed palette (validated categorical order: blue, orange, aqua).
BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
INK, INK2, GRID = "#0b0b0b", "#52514e", "#e4e3df"


def _style(ax, title=None, xlabel=None, ylabel=None):
    ax.set_facecolor("white")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("left", "bottom"):
        ax.spines[s].set_color(GRID)
    ax.tick_params(colors=INK2, labelsize=8)
    ax.grid(True, color=GRID, linewidth=0.6)
    ax.set_axisbelow(True)
    if title:
        ax.set_title(title, fontsize=10, color=INK, loc="left")
    if xlabel:
        ax.set_xlabel(xlabel, fontsize=9, color=INK2)
    if ylabel:
        ax.set_ylabel(ylabel, fontsize=9, color=INK2)


def _png(fig) -> str:
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=110, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def fig_reliability(P_before, P_after, y, label_after="calibrated") -> str:
    fig, axes = plt.subplots(1, 2, figsize=(8, 3.6), sharey=True)
    for ax, P, name, color in ((axes[0], P_before, "uncalibrated", ORANGE), (axes[1], P_after, label_after, BLUE)):
        bconf, bacc, cnt, edges = reliability_bins(P, y, 10)
        m = cnt > 0
        centers = (edges[:-1] + edges[1:]) / 2
        ax.bar(centers[m], bacc[m], width=0.09, color=color, alpha=0.85, label="accuracy in bin")
        ax.plot([0, 1], [0, 1], color=INK2, linewidth=1, linestyle="--", label="perfect calibration")
        ax.plot(bconf[m], bacc[m], "o", color=INK, markersize=3, label="(mean conf, acc)")
        _style(ax, f"Reliability: {name}", "confidence (max prob)", "accuracy" if ax is axes[0] else None)
        ax.set_xlim(0, 1); ax.set_ylim(0, 1)
    axes[0].legend(fontsize=7, frameon=False, loc="upper left")
    return _png(fig)


def fig_class_coverage(conformal: dict, classes, alpha) -> str:
    K = len(classes)
    fig, ax = plt.subplots(figsize=(max(5, 0.5 * K + 3), 3.4))
    x = np.arange(K)
    w = 0.38
    for off, key, color, name in ((-w / 2, "aps", ORANGE, "APS (marginal)"),
                                  (w / 2, "aps_mondrian", BLUE, "APS (Mondrian / per-class)")):
        cov = [conformal[key]["per_class_coverage"][k] for k in range(K)]
        ax.bar(x + off, cov, width=w - 0.03, color=color, label=name)
    ax.axhline(1 - alpha, color=INK, linewidth=1, linestyle="--", label=f"target {1 - alpha:.2f}")
    ax.set_xticks(x, [str(c) for c in classes], rotation=45 if K > 6 else 0, ha="right" if K > 6 else "center")
    lo = min(min(conformal[k]["per_class_coverage"].values()) for k in ("aps", "aps_mondrian"))
    ax.set_ylim(max(0, lo - 0.1), 1.01)
    _style(ax, "Coverage by true class (test split)", None, "coverage")
    ax.legend(fontsize=7, frameon=False, loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=3)
    return _png(fig)


def fig_set_sizes(conformal: dict) -> str:
    fig, ax = plt.subplots(figsize=(5, 3.2))
    for key, color, name, off in (("lac", AQUA, "LAC", -0.27), ("aps", ORANGE, "APS", 0.0),
                                  ("aps_mondrian", BLUE, "APS Mondrian", 0.27)):
        h = np.asarray(conformal[key]["size_hist"], float)
        ax.bar(np.arange(len(h)) + off, h / h.sum(), width=0.25, color=color, label=name)
    _style(ax, "Prediction-set size distribution", "set size", "fraction of test points")
    ax.legend(fontsize=7, frameon=False)
    return _png(fig)


def fig_risk_coverage(sel: dict) -> str:
    fig, ax = plt.subplots(figsize=(5, 3.4))
    cov = np.asarray(sel["test_curve"]["coverage"]); risk = np.asarray(sel["test_curve"]["risk"])
    ax.plot(cov, risk, color=BLUE, linewidth=2, label="test risk-coverage")
    ax.axhline(sel["target_risk"], color=INK2, linestyle="--", linewidth=1, label=f"target risk {sel['target_risk']}")
    if sel.get("threshold") is not None:
        ax.plot([sel["test_coverage"]], [sel["test_selective_risk"]], "o", color=ORANGE, markersize=8,
                label=f"chosen threshold {sel['threshold']:.3f}")
    _style(ax, "Selective risk vs coverage", "coverage (fraction answered)", "error among answered")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, max(sel["target_risk"], float(risk.max())) * 1.15)
    ax.legend(fontsize=7, frameon=False, loc="upper center", bbox_to_anchor=(0.5, -0.18), ncol=3)
    return _png(fig)


def fig_certified(cert: dict) -> str:
    fig, ax = plt.subplots(figsize=(5, 3.4))
    ax.plot(cert["radii"], cert["certified_accuracy"], color=BLUE, linewidth=2)
    _style(ax, "Certified accuracy (logistic regression)", "L2 radius in model feature space",
           "correct and certified")
    ax.set_ylim(0, 1)
    return _png(fig)


def _table(rows: list[dict], cols=None, fmt="{:.3f}") -> str:
    if not rows:
        return "<p>(none)</p>"
    cols = cols or list(rows[0])
    head = "".join(f"<th>{html.escape(str(c))}</th>" for c in cols)
    body = ""
    for r in rows:
        cells = []
        for c in cols:
            v = r.get(c, "")
            if isinstance(v, float):
                v = fmt.format(v)
            cells.append(f"<td>{html.escape(str(v))}</td>")
        body += "<tr>" + "".join(cells) + "</tr>"
    return f"<table><thead><tr>{head}</tr></thead><tbody>{body}</tbody></table>"


CSS = """
body{font-family:system-ui,-apple-system,Segoe UI,sans-serif;max-width:980px;margin:24px auto;padding:0 16px;
color:#0b0b0b;background:#fcfcfb;line-height:1.5}
h1{font-size:1.6rem;margin-bottom:0} h2{font-size:1.2rem;margin-top:2rem;border-bottom:1px solid #e4e3df}
.muted{color:#52514e;font-size:.9rem} img{max-width:100%;height:auto;border:1px solid #e4e3df;border-radius:6px;
background:white;margin:6px 0} table{border-collapse:collapse;font-size:.85rem;margin:8px 0;display:block;
overflow-x:auto} th,td{padding:4px 10px;border-bottom:1px solid #e4e3df;text-align:right}
th{background:#f3f2ee} td:first-child,th:first-child{text-align:left}
.kpi{display:inline-block;margin:6px 18px 6px 0} .kpi b{display:block;font-size:1.3rem}
.alarm{color:#b3261e;font-weight:600} .ok{color:#00703c;font-weight:600}
"""


def render_html(res: dict, bundle, drift_res: dict | None = None) -> str:
    """Build the report HTML from ``pipeline.evaluate`` output (and optional drift results)."""
    classes = res["classes"]
    alpha = res["alpha"]
    m = res["test_metrics"]
    parts = [f"<html><head><meta charset='utf-8'><meta name='viewport' content='width=device-width,initial-scale=1'>"
             f"<title>clfkit report</title><style>{CSS}</style></head><body>"]
    parts.append("<h1>clfkit report</h1>")
    parts.append(f"<p class='muted'>model <b>{res['model']}</b> &middot; {len(classes)} classes &middot; "
                 f"split sizes {res['split_sizes']} &middot; calibration: {res['calibration_used']}</p>")
    parts.append("".join(f"<span class='kpi'><b>{v}</b><span class='muted'>{k}</span></span>" for k, v in [
        ("test accuracy", f"{m['accuracy']:.3f}"), ("test ECE", f"{m['ece']:.3f}"),
        (f"APS coverage (target {1 - alpha:.2f})", f"{res['conformal']['aps']['coverage']:.3f}"),
        ("APS avg set size", f"{res['conformal']['aps']['avg_set_size']:.2f}"),
    ]))

    parts.append("<h2>1. Model comparison</h2><p class='muted'>Internal 80/20 split of the training rows; "
                 "sorted by validation log-loss (a proper scoring rule).</p>")
    parts.append(_table(res["comparison"]))

    parts.append("<h2>2. Calibration</h2><p class='muted'>Each calibrator is fitted on <i>calib_fit</i> "
                 "and evaluated on <i>test</i>. ECE: binned |accuracy - confidence| (15 bins). "
                 "NLL and Brier are proper scores (lower is better).</p>")
    parts.append(_table(res["calibration"]))
    parts.append(f"<img src='data:image/png;base64,{fig_reliability(bundle.P_test_raw, bundle.P_test, bundle.y_test, res['calibration_used'])}'>")

    c = res["conformal"]
    parts.append("<h2>3. Conformal prediction sets</h2>")
    parts.append(f"<p class='muted'>Thresholds from <i>calib_conf</i> (n={res['conformal_n_cal']}); "
                 f"coverage measured on <i>test</i>. Guarantee: marginal coverage &ge; {1 - alpha:.2f} "
                 f"(and &le; {res['conformal_upper']:.3f} for continuous scores) on average over calibration "
                 f"draws, if data are exchangeable. Per-class coverage is <b>not</b> guaranteed by the marginal "
                 f"method; Mondrian conformal fixes this per class at the cost of larger sets.</p>")
    rows = [{"method": k, "coverage": v["coverage"], "avg_set_size": v["avg_set_size"],
             "singletons": v["frac_singleton"], "empty": v["frac_empty"],
             "worst_class_cov": v["worst_class_coverage"]} for k, v in c.items()]
    parts.append(_table(rows))
    parts.append(f"<img src='data:image/png;base64,{fig_class_coverage(c, classes, alpha)}'>")
    parts.append(f"<img src='data:image/png;base64,{fig_set_sizes(c)}'>")

    s = res["selective"]
    parts.append("<h2>4. Abstention (selective classification)</h2>")
    if s["threshold"] is None:
        parts.append(f"<p class='alarm'>No threshold certifies selective risk &le; {s['target_risk']} "
                     f"with confidence {1 - s['delta']:.2f} on n={s['n_holdout']} held-out points. "
                     "Collect more calibration data or relax the target.</p>")
    else:
        cr = s["certified"]
        parts.append(
            f"<p>Answer only if confidence &ge; <b>{s['threshold']:.4f}</b>. On <i>calib_conf</i>: "
            f"{cr['n_selected']} answered, {cr['errors']} wrong, simultaneous upper bound on selective risk "
            f"{cr['risk_upper_bound']:.4f} &le; target {s['target_risk']} "
            f"(Clopper-Pearson, Bonferroni over {s['grid_size']} fixed thresholds, &delta;={s['delta']}).</p>"
            f"<p>On <i>test</i>: coverage {s['test_coverage']:.3f}, selective risk {s['test_selective_risk']:.4f}. "
            f"AURC {s['test_aurc']:.4f}.</p>"
            "<p class='muted'>Guaranteed: with prob. &ge; 1-&delta; over the held-out draw, the true error rate "
            "among answered inputs is &le; target, for future data from the same distribution. Not guaranteed: "
            "per-input correctness, the realised rate on a small batch, the fraction answered, or anything under drift.</p>")
    parts.append(f"<img src='data:image/png;base64,{fig_risk_coverage(s)}'>")

    if "certified" in res:
        ce = res["certified"]
        parts.append("<h2>5. Robustness</h2>")
        parts.append(
            f"<p>Exact L2 certificate for the <b>logistic-regression</b> baseline in its {html.escape(ce['space'])}: "
            f"no perturbation of norm below the radius can change its predicted label. "
            f"Clean accuracy {ce['clean_accuracy']:.3f}; median certified radius {ce['median_radius']:.3f}.</p>")
        parts.append(f"<img src='data:image/png;base64,{fig_certified(ce)}'>")
        st = res["stability"]
        parts.append(f"<p>Empirical stability of the final model (<b>{st['model']}</b>) under random noise "
                     "(numeric: Gaussian with sd = &sigma; &times; feature sd; text: each token dropped with "
                     "prob. &sigma;). Random noise is much weaker than adversarial noise: this is a smoke test, "
                     "not a certificate.</p>")
        parts.append(_table(st["rows"]))

    if drift_res is not None:
        parts.append(render_drift_section(drift_res))

    parts.append("<p class='muted' style='margin-top:3rem'>Generated by clfkit. Raw numbers: summary.json.</p>")
    parts.append("</body></html>")
    return "\n".join(parts)


def render_drift_section(d: dict) -> str:
    sm = d["summary"]
    status = "<span class='alarm'>DRIFT DETECTED</span>" if sm["drift"] else "<span class='ok'>no drift detected</span>"
    out = [f"<h2>6. Drift check</h2><p>{status} (level {sm['level']}; reference n={d['n_ref']}, new n={d['n_new']}). "
           f"Alarms: {html.escape(json.dumps(sm['alarms']))}</p>"]
    out.append(f"<p>Score drift (KS on confidence): stat {d['scores']['statistic']:.3f}, p={d['scores']['p']:.2e}; "
               f"mean confidence {d['scores']['mean_conf_ref']:.3f} &rarr; {d['scores']['mean_conf_new']:.3f}.<br>"
               f"Classifier two-sample test: held-out accuracy {d['c2st']['accuracy']:.3f} "
               f"(chance 0.5), p={d['c2st']['p']:.2e}.</p>")
    out.append(_table(d["features"][:15], ["feature", "test", "statistic", "p", "p_holm"], fmt="{:.3g}"))
    ls = d["label_shift"]
    rows = [{"class": c, "reference": p, "estimated_new (BBSE)": q, "predicted_new": mu}
            for c, p, q, mu in zip(ls["classes"], ls["p_ref"], ls["q_est"], ls["pred_dist_new"])]
    out.append(f"<p>Label-shift estimate (valid only if p(x|y) is unchanged; confusion-matrix condition "
               f"number {ls['condition_number']:.1f}):</p>")
    out.append(_table(rows))
    return "\n".join(out)


def write_report(res: dict, bundle, path: str | Path, drift_res: dict | None = None) -> Path:
    path = Path(path)
    path.write_text(render_html(res, bundle, drift_res), encoding="utf-8")
    return path