"""Render schema-v1 performance results without recomputing measurements."""

from __future__ import annotations

import argparse
import html
import json
from pathlib import Path
import re
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from safety import SecretScanner  # noqa: E402

from records import COLLECTION, DID, LIVE_COLLECTION  # noqa: E402

SCANNER = SecretScanner()


def write_file(path, content):
    SCANNER.check(content, label=str(path))
    if isinstance(content, bytes):
        Path(path).write_bytes(content)
    else:
        Path(path).write_text(content)


def read_result(path):
    content = Path(path).read_text()
    SCANNER.check(content, label=str(path))
    result = json.loads(content)
    if result["schema"] != 1 or result["mode"] not in ("live", "replay"):
        raise ValueError(f"unsupported performance result: {path}")
    if not re.fullmatch(r"[0-9a-f]{40}", result["commit"]):
        raise ValueError("commit must be a full lowercase SHA")
    return result


def esc(value):
    return html.escape(str(value), quote=True)


def number(value):
    return "—" if value is None else f"{value:,.2f}".rstrip("0").rstrip(".")


def table(headers, rows):
    head = "".join(f"<th>{esc(h)}</th>" for h in headers)
    body = "".join(
        "<tr>"
        + "".join(
            f'<td data-label="{esc(h)}">{esc(v)}</td>'
            for h, v in zip(headers, row, strict=True)
        )
        + "</tr>"
        for row in rows
    )
    return f"<table><thead><tr>{head}</tr></thead><tbody>{body}</tbody></table>"


ECHARTS = "https://cdn.jsdelivr.net/npm/echarts@6.1.0/dist/echarts.min.js"
# sha384 of dist/echarts.min.js from the npm tarball, checked against the
# registry's sha512 and against jsDelivr's copy (byte-identical).
ECHARTS_SRI = "sha384-C2iskrW/uPW46KzOjrvJIQo4YkV8lkD+QS0CrDN18IIPIpT/g2USu8bTP3nvmIAD"


def summary(result):
    """The numbers pages compare across runs, small enough for every record."""
    return {
        "gated": result["gated"],
        "scenarios": {
            name: {
                "memory": {m: v["median"] for m, v in scenario["memory"].items()},
                "wire": scenario["wire"],
            }
            for name, scenario in result["scenarios"].items()
        },
    }


TILES = (
    ("boot_tree_uss", "Memory at boot"),
    ("idle_session_tree_uss", "Memory per idle session"),
    ("retained_turn_beam", "Memory kept per turn"),
    ("turn_peak_tree_uss", "Peak during a turn"),
    ("request_bytes_per_turn", "Sent to the model per turn"),
    ("attach_bytes", "Sent to open a long session"),
)


def size(value):
    """Bytes in decimal units with three significant figures."""
    if value is None:
        return "-"
    for unit, scale in (("GB", 1e9), ("MB", 1e6), ("KB", 1e3)):
        if abs(value) >= scale:
            return f"{value / scale:.3g} {unit}"
    return f"{value:.0f} B"


def median_observation(scenario):
    """The run whose final memory is the median, so its timeline is typical."""
    runs = [o for o in scenario.get("observations", []) if o.get("checkpoints")]
    if not runs:
        return {"checkpoints": [], "turns": []}
    finals = sorted(
        range(len(runs)), key=lambda i: private(runs[i]["checkpoints"][-1]) or 0
    )
    return runs[finals[len(finals) // 2]]


def private(snapshot):
    """A snapshot's private memory (USS), or its PSS from results that predate USS."""
    return snapshot.get("tree_uss", snapshot.get("tree_pss"))


def median(value):
    """A plain number, or the median of a {median, mad, samples} distribution."""
    return value.get("median") if isinstance(value, dict) else value


def pool_report(pool):
    """The pool scenario's warm pool and the first and adopted kernel."""
    if not pool:
        return None
    kernels = {
        name: {
            "kinds": sorted(set(pool[name].get("kinds", []))),
            "uss": median(pool[name].get("uss")),
            "pss": median(pool[name].get("pss")),
        }
        for name in ("first_kernel", "adopted_kernel")
        if name in pool
    }
    return {"warmPss": median(pool.get("warm_pss")), **kernels}


def timeline(scenario):
    """What the scenario charts draw: settled memory and per-request bytes."""
    observation = median_observation(scenario)
    return {
        "checkpoints": [
            {
                "label": c["label"],
                "count": c["count"],
                "memory": private(c),
            }
            for c in observation["checkpoints"]
        ],
        # Keyed by turn number: checkpoints are sparse, and not every turn
        # carries a peak.
        "peaks": {
            str(t["turn"]): t.get("peak_tree_uss", t.get("peak_tree_pss"))
            for t in observation.get("turns", [])
            if "peak_tree_uss" in t or "peak_tree_pss" in t
        },
        "requests": [
            {"sent": r["sent"], "received": r["received"]}
            for r in scenario.get("requests", [])
        ],
        "wire": scenario.get("wire", {}),
        "memory": {m: v["median"] for m, v in scenario.get("memory", {}).items()},
        "wallMs": scenario.get("wall_ms", {}).get("median"),
        "pool": pool_report(scenario.get("pool")),
        "runs": scenario.get("runs"),
    }


def page_data(result):
    """Everything the page script needs about this run, without raw samples."""
    recordings = {
        name: {
            "recordedAt": scenario["recordings"]["recorded_at"],
            "model": scenario["recordings"]["model"],
            "stale": bool(scenario["recordings"]["stale_turns"]),
        }
        for name, scenario in result["scenarios"].items()
        if scenario.get("recordings")
    }
    return {
        "commit": result["commit"],
        "baseCommit": result["base_commit"],
        "ref": result["ref"],
        "event": result["event"],
        "verdict": result["verdict"],
        "gated": {
            name: {
                **{
                    k: m[k] for k in ("value", "noise", "baseline", "budget", "verdict")
                },
                "singleSample": bool(m.get("single_sample")),
            }
            for name, m in result["gated"].items()
        },
        "scenarios": {n: timeline(s) for n, s in result["scenarios"].items()},
        "recordings": recordings,
        "warmup": warmup(result),
    }


def warmup(result):
    """The long scenario's one-time warm-up, when the result carries it."""
    measured = result["scenarios"].get("long", {}).get("warmup_retained")
    if not measured:
        return None
    private = measured.get("tree_uss") or measured["tree_pss"]
    return {"turns": measured["turns"], "memory": private["median"]}


def fallback(result):
    """The gate's numbers for readers without JavaScript."""
    rows = "".join(
        f"<tr><td>{esc(label)}</td><td>{esc(size(result['gated'][key]['value']))}</td>"
        f"<td>{esc(result['gated'][key]['verdict'])}</td></tr>"
        for key, label in TILES
        if key in result["gated"]
    )
    return (
        f"<noscript><p>Gate {esc(result['verdict'])}.</p>"
        f"<table><tbody>{rows}</tbody></table></noscript>"
    )


def page(title, body, own=None):
    scripts = ""
    if own is not None:
        payload = json.dumps(own, separators=(",", ":")).replace("<", "\\u003c")
        scripts = (
            f'<script type="application/json" id="run">{payload}</script>'
            f'<script src="{ECHARTS}" integrity="{ECHARTS_SRI}" crossorigin="anonymous"></script>'
            '<script src="app.js"></script>'
        )
    return f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>{esc(title)} · albedo perf</title><link rel="stylesheet" href="style.css">
</head><body>{body}{scripts}</body></html>"""


SKELETON = """<header class="bar">
  <span class="chip" id="chip"></span>
  <p class="verdict" id="verdict">Loading.</p>
  <div class="meta">
    <span class="hash" id="commit"></span><span id="branch"></span><span id="when"></span>
    <label class="compare">vs <select id="compare" disabled></select></label>
  </div>
</header>
<main>
  <section class="tiles" id="tiles" aria-label="Headline numbers"></section>
  <section id="moved" aria-labelledby="moved-title"><h2 id="moved-title">What moved</h2><div id="moved-body"></div></section>
  <section aria-labelledby="scenarios-title">
    <div class="section-head"><h2 id="scenarios-title">Scenarios</h2>
      <div class="segmented" id="mode" role="group" aria-label="Data source"></div></div>
    <p class="note" id="mode-note"></p>
    <div class="scenarios" id="scenarios"></div>
  </section>
  <section aria-labelledby="history-title"><div class="section-head"><h2 id="history-title">History on main</h2>
    <div class="tabs" id="history-tabs" role="tablist"></div></div><div id="history"></div></section>
  <footer id="footer"></footer>
</main>"""


def build_run_site(directory, result, own):
    """Write one run's immutable site: its page, its data and its assets."""
    directory = Path(directory)
    (directory / "data").mkdir(parents=True, exist_ok=True)
    SCANNER.write_json(directory / "data/run.json", result)
    own = {**own, "liveCollection": LIVE_COLLECTION, **page_data(result)}
    body = SKELETON.replace("</main>", fallback(result) + "</main>")
    write_file(
        directory / "index.html",
        page(f"{result['commit'][:7]} · {result['ref']}", body, own),
    )
    for asset in ("style.css", "app.js"):
        write_file(directory / asset, Path(__file__).with_name(asset).read_bytes())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Render one run's site for a local preview."
    )
    parser.add_argument("site", type=Path)
    parser.add_argument("result", type=Path)
    parser.add_argument("--did", default=DID, help="repo to read runs from")
    parser.add_argument("--pds", help="PDS used when the DID does not resolve")
    parser.add_argument("--created-at", help="when the run was published (ISO 8601)")
    args = parser.parse_args()
    result = read_result(args.result)
    build_run_site(
        args.site,
        result,
        {
            "did": args.did,
            "collection": COLLECTION,
            "rkey": result["commit"],
            "pds": args.pds,
            "createdAt": args.created_at,
        },
    )