"""Guards for ../widget.html: its data block, its size, that its random run and sentences are the ones speculative_sim.py writes, and what a reader sees."""
import json
import re
import subprocess

import pytest

import speculative_sim as sim
from embed_widget import CONFIG_BLOCK, WIDGET, config, embed

HTML = WIDGET.read_text(encoding="utf-8")
ROOT = HTML[HTML.index('<div class="sdw-root">'):HTML.index("<script")]
SLOWER = "On average, slower than plain decoding: the draft passes cost more than the kept guesses save."
GRID = [(pct, k, d) for pct in range(0, 101, 5) for k in range(1, 9) for d in range(len(sim.DRAFT_STOPS_MS))]


def settings_view(pct, k, draft_index, seed):
    return sim.view(pct, k, sim.DRAFT_STOPS_MS[draft_index], seed)


# Reads the model out of the widget between its marker comments, runs it with Node, and compares the random numbers, the rounds and every
# sentence with the ones Python wrote into the fixture (the two times as numbers, the rest as exact text).
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const cfg = JSON.parse(html.match(/<script type="application\/json" id="sdw-config">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the model ('), html.indexOf('/* ---------- display'));
const api = new Function('cfg', "'use strict';" + js + '; return {stream, planRun, view, expectedTokens};')(cfg);
const bad = [];
let checks = 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
const near = (a, b) => Math.abs(a - b) <= 1e-12 * Math.abs(b);
fx.streams.forEach((s) => {
  const draw = api.stream(s.seed);
  const got = s.draws.map(() => draw());
  checks++;
  if (!same(got, s.draws)) bad.push(['stream', s.seed, got.slice(0, 3), s.draws.slice(0, 3)]);
});
fx.cases.forEach((c) => {
  const v = api.view(c.pct, c.k, cfg.draftStops[c.d], c.seed);
  checks += 3;
  if (!same(v.rounds, c.view.rounds)) bad.push(['rounds', c.pct, c.k, c.seed, v.rounds, c.view.rounds]);
  if (!['plain', 'spec', 'verdict', 'average', 'draft'].every((key) => v[key] === c.view[key])) bad.push(['text', c.pct, c.k, c.d, v, c.view]);
  if (!near(v.plainMs, c.view.plainMs) || !near(v.roundMs, c.view.roundMs)) bad.push(['times', c.pct, c.k, c.d, v.plainMs, v.roundMs]);
});
fx.expected.forEach(([a, k, e]) => { checks++; if (!near(api.expectedTokens(a, k), e)) bad.push(['expected', a, k, api.expectedTokens(a, k), e]); });
console.log(JSON.stringify({ checks, bad: bad.slice(0, 4) }));
process.exit(bad.length ? 1 : 0);
"""

# Runs the page script against a small stand-in for the DOM, with its timer and reduced-motion setting replaced, moves the sliders, presses the buttons
# and reads back what a reader sees: the words in each lane, the sentences, what is spoken, what Replay shows frame by frame.
PAGE = r"""
import fs from 'node:fs';
const [widget, fixture] = process.argv.slice(2);
const html = fs.readFileSync(widget, 'utf8');
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const config = html.match(/<script type="application\/json" id="sdw-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, value: '', hidden: false, attrs: {}, text: '' };
  Object.defineProperty(e, 'textContent', { get: () => e.text, set: (v) => { e.text = v; e.children.length = 0; } });
  e.setAttribute = (k, v) => { e.attrs[k] = v; };
  e.addEventListener = (type, f) => (e.listeners[type] = e.listeners[type] || []).push(f);
  e.appendChild = (c) => (e.children.push(c), c);
  return e;
}
function open(reduced) {
  const page = { root: el('html'), messages: [], parent: { other: true }, timers: [], cleared: [] };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['sdw-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f), matchMedia: (q) => ({ matches: reduced && q === '(prefers-reduced-motion: reduce)' }) };
  const setTimer = (f, ms) => (page.timers.push({ f, ms }), page.timers.length);
  new Function('document', 'window', 'parent', 'setInterval', 'clearInterval', code)(doc, win, page.parent, setTimer, (id) => page.cleared.push(id));
  return page;
}
const fire = (e, type) => (e.listeners[type || 'click'] || []).forEach((f) => f({ target: e }));
const bad = [];
const want = (what, got, expected) => { if (JSON.stringify(got) !== JSON.stringify(expected)) bad.push([what, got, expected]); };
const guesses = (k) => k + (k === 1 ? ' guess' : ' guesses');
const speech = (pct, k, v) => pct + '% of guesses kept, ' + guesses(k) + ' per round. ' + v.spec + (v.verdict ? ' ' + v.verdict : '');

function passes(page) { return page['sdw-spec'].children; }
function look(page, v, what, pct, k) {              // compares the page with one expected view
  ['plain', 'spec', 'verdict', 'average'].forEach((key) => want(what + ': ' + key, page['sdw-' + (key === 'plain' ? 'plain-sum' : key === 'spec' ? 'spec-sum' : key)].textContent, v[key]));
  want(what + ': percent label', page['sdw-pct-n'].textContent, pct + '%');
  want(what + ': guesses label', page['sdw-k-n'].textContent, String(k));
  want(what + ': draft label', page['sdw-draft-n'].textContent, v.draft);
  want(what + ': slider readouts', [page['sdw-pct'].attrs['aria-valuetext'], page['sdw-k'].attrs['aria-valuetext'], page['sdw-draft'].attrs['aria-valuetext']],
       [pct + ' percent', guesses(k), v.draft + ' of a big pass']);
  const plain = page['sdw-plain'].children;
  want(what + ': plain lane words', plain.map((c) => [c.tag, c.attrs['data-kind'], c.textContent]), fx.reply.map((w) => ['strong', 'own', w]));
  want(what + ': pass numbers', passes(page).map((g) => g.attrs['data-pass']), v.rounds.map((r, j) => String(j + 1)));
  const expected = v.rounds.map((r) => {
    const chips = [];
    for (let i = 0; i < r[1]; i++) chips.push(['span', 'kept', fx.reply[r[0] + i]]);
    if (r[2]) chips.push(['del', 'rejected', fx.decoys[r[0] + r[1]]]);
    if (r[3]) chips.push(['strong', 'own', fx.reply[r[0] + r[1]]]);
    return chips;
  });
  want(what + ': passes', passes(page).map((g) => g.children.map((c) => [c.tag, c.attrs['data-kind'], c.textContent])), expected);
  const spelled = [];                                // the kept guesses and the big model's own words, in order, are the reply: nothing lost, nothing added
  passes(page).forEach((g) => g.children.forEach((c) => { if (c.attrs['data-kind'] !== 'rejected') spelled.push(c.textContent); }));
  want(what + ': the reply is spelled out', spelled, fx.reply);
  want(what + ': all shown', plain.concat(passes(page)).map((c) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);
}

const first = open(false);
const opening = fx.opening;
look(first, opening.view, 'opening state', opening.pct, opening.k);
want('slider values', [first['sdw-pct'].value, first['sdw-k'].value, first['sdw-draft'].value], [String(opening.pct), String(opening.k), String(fx.defaultDraft)]);
want('nothing spoken yet', first['sdw-live'].textContent, '');
want('replay shown', first['sdw-replay'].hidden, false);
want('nothing animating at the start', first.timers.length, 0);

// every combination: dragging stays silent, releasing speaks
fx.combos.forEach((c) => {
  const page = open(false);
  page['sdw-pct'].value = String(c.pct); page['sdw-k'].value = String(c.k); page['sdw-draft'].value = String(c.d);
  fire(page['sdw-pct'], 'input');
  look(page, c.view, 'setting ' + [c.pct, c.k, c.d], c.pct, c.k);
  want('dragging stays silent at ' + [c.pct, c.k, c.d], page['sdw-live'].textContent, '');
  fire(page['sdw-k'], 'change');
  want('release speaks at ' + [c.pct, c.k, c.d], page['sdw-live'].textContent, speech(c.pct, c.k, c.view));
});

// New run draws the next seed and speaks
fx.newRuns.forEach((c) => {
  const page = open(false);
  page['sdw-pct'].value = String(c.pct); page['sdw-k'].value = String(c.k); page['sdw-draft'].value = String(c.d);
  fire(page['sdw-new']);
  look(page, c.view, 'new run at ' + [c.pct, c.k, c.d], c.pct, c.k);
  want('new run speaks', page['sdw-live'].textContent, speech(c.pct, c.k, c.view));
});

// Replay: everything hidden, then each pass appears when the clock passes its end; the last frame shows all and announces once
const replaying = open(false);
fire(replaying['sdw-replay']);
const lanes = () => replaying['sdw-plain'].children.map((c) => [c, (replaying['sdw-plain'].children.indexOf(c) + 1) * opening.view.plainMs])
  .concat(passes(replaying).map((g, j) => [g, (j + 1) * opening.view.roundMs]));
const end = Math.max(fx.reply.length * opening.view.plainMs, opening.view.rounds.length * opening.view.roundMs);
want('replay starts hidden', lanes().map(([c]) => c.attrs['data-wait']).filter((w) => w !== 'yes'), []);
want('one timer at 40 ms', replaying.timers.map((t) => t.ms), [40]);
const tick = replaying.timers[0].f;
let frame = 0;
[1, 25, 50, 75, 99].forEach((stop) => {
  while (frame < stop) { tick(); frame++; }
  const now = end * frame / 100;
  want('replay frame ' + frame, lanes().map(([c, at]) => c.attrs['data-wait'] === (at <= now ? 'no' : 'yes')).filter((ok) => !ok), []);
  want('silent while replaying at ' + frame, replaying['sdw-live'].textContent, '');
});
const shownAt = (stop) => {                          // the shorter lane is finished while the plain one is still writing
  const now = end * stop / 100;
  return [opening.view.rounds.length * opening.view.roundMs <= now, fx.reply.length * opening.view.plainMs <= now];
};
want('the draft lane finishes before the plain one', shownAt(75), [true, false]);
tick(); frame++;
want('last frame shows everything', lanes().map(([c]) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);
want('last frame announces', replaying['sdw-live'].textContent, 'Replay finished. ' + speech(opening.pct, opening.k, opening.view));
want('timer stopped', replaying.cleared.includes(1), true);

// moving a slider during a replay stops it and shows the finished state
const interrupted = open(false);
fire(interrupted['sdw-replay']);
interrupted['sdw-pct'].value = '50';
fire(interrupted['sdw-pct'], 'input');
want('slider stops the replay', interrupted.cleared.includes(1), true);
want('slider shows the finished state', interrupted['sdw-plain'].children.concat(passes(interrupted)).map((c) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);

// reduced motion: no Replay button, finished state, nothing animating
const calm = open(true);
want('reduced motion hides Replay', calm['sdw-replay'].hidden, true);
want('reduced motion starts nothing', calm.timers.length, 0);
look(calm, opening.view, 'reduced motion opening', opening.pct, opening.k);

// the host page theme: only the parent window may set it, and only to light or dark
const themed = open(false);
const send = (data, source) => themed.messages.forEach((f) => f({ data, source }));
send({ type: 'spicyshark:theme', theme: 'dark' }, { other: false });
send({ type: 'spicyshark:theme', theme: 'purple' }, themed.parent);
send({ type: 'other', theme: 'dark' }, themed.parent);
send(null, themed.parent);
want('ignored messages', themed.root.attrs, {});
send({ type: 'spicyshark:theme', theme: 'dark' }, themed.parent);
want('theme from the parent', themed.root.attrs, { 'data-theme': 'dark' });
send({ type: 'spicyshark:theme', theme: 'light' }, themed.parent);
want('theme back to light', themed.root.attrs, { 'data-theme': 'light' });

console.log(JSON.stringify({ combos: fx.combos.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def run_node(tmp_path, script, fixture, args):
    (tmp_path / "fixture.json").write_text(json.dumps(fixture), encoding="utf-8")
    (tmp_path / "check.mjs").write_text(script, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "check.mjs"), *args(tmp_path)], capture_output=True, text=True, timeout=120)
    assert run.returncode == 0, run.stdout + run.stderr
    return json.loads(run.stdout)


def test_widget_data_block_matches_the_python_constants():
    """Guards a constant edit that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_embed_rewrites_a_widget_whose_data_block_drifted():
    """Guards embed() skipping the block: a hand-edited parameter count or opening draft stop must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"defaultDraft": 2', '"defaultDraft": 0', 1)
    assert drifted != HTML and embed(drifted) == HTML
    assert json.loads(CONFIG_BLOCK.search(HTML).group(2)) == json.loads(json.dumps(config()))


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: the parity run covers only the model, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    assert "$('sdw-' + name)" in script                   # the four sentences are looked up by suffix; they are added below
    lookups = set(re.findall(r"'(sdw-[\w-]+)'", script)) | {f"sdw-{n}" for n in ("plain-sum", "spec-sum", "verdict", "average")}
    missing = sorted(i for i in lookups if i != "sdw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="sdw-config"' in HTML


def test_widget_is_small_and_asks_for_nothing_online():
    assert len(HTML.splitlines()) <= 300 and len(HTML.encode("utf-8")) < 40 * 1024
    page = CONFIG_BLOCK.sub("", HTML)
    for banned in ("http://", "https://", "fetch(", "XMLHttpRequest", "WebSocket", "sendBeacon", "import("):
        assert banned not in page, banned
    assert " src=" not in page and "@import" not in page


def test_widget_has_no_form_element():
    """Guards a dead button: the lesson page loads widget.html in a sandboxed frame, where a form cannot submit."""
    assert "<form" not in HTML and 'type="submit"' not in HTML


def test_widget_shows_at_most_five_controls_and_hides_the_rest():
    """Guards the first screen at 390 px: two sliders and Replay stay in view; the draft slider and New run sit in a closed details element."""
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert len(re.findall(r"<(?:select|button|input|textarea)\b", outside)) <= 5
    more = re.search(r"<details><summary>More settings</summary>(.*?)</details>", ROOT, re.S).group(1)
    assert 'id="sdw-draft"' in more and 'id="sdw-new"' in more and "<details open" not in ROOT


def test_widget_keeps_the_site_theme_listener_and_palette():
    assert "d.type !== 'spicyshark:theme'" in HTML and "e.source !== parent" in HTML and "parent === window" in HTML
    for pair in ("--bg: light-dark(#ffffff, #0a0f17)", "--panel: light-dark(#ebf0f6, #131b28)", "--ink: light-dark(#0a1626, #f4f7fb)",
                 "--mute: light-dark(#44566b, #a3b1c4)", "--ctl: light-dark(#66768a, #7a8a9f)", "--ink-on: light-dark(#f6f8fb, #05070c)",
                 "--green: light-dark(#00855F, #009E73)", "--orange: light-dark(#b36200, #E69F00)", ":root[data-theme=light]", ":root[data-theme=dark]"):
        assert pair in HTML, pair
    for banned in ("#0072B2", "#56B4E9", "#D55E00"):
        assert banned.lower() not in HTML.lower()
    assert not set(re.findall(r"#[0-9a-fA-F]{3,8}\b", HTML)) - {"#ffffff", "#0a0f17", "#ebf0f6", "#131b28", "#0a1626", "#f4f7fb", "#44566b", "#a3b1c4",
                                                               "#66768a", "#7a8a9f", "#f6f8fb", "#05070c", "#00855F", "#009E73", "#b36200", "#E69F00"}


def test_the_visible_notes_quote_what_the_code_computes():
    """Guards a stale sentence: the More settings note names the 1.24-billion-weight draft model, its 1.2 ms pass and the 7.88 ms big pass, and the
    left-out list says about 2.5 GB. Also the words that keep a best case from reading as a measurement, and the key that names each kind of chip."""
    view = settings_view(80, 4, sim.DEFAULT_DRAFT_STOP, sim.OPENING["seed"])
    assert f"{sim.DRAFT_PARAMS / 1e9:.2f}-billion-weight" in HTML and view["draft"].startswith("1.2 ms") and "The opening stop, 1.2 ms," in HTML
    assert view["plain"].endswith("7.88 ms per token.") and "takes 7.88 ms" in HTML and f"about {sim.DRAFT_PARAMS * 2 / 1e9:.1f} GB" in HTML
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert "A model at the chip&rsquo;s peak speed, not a measurement." in outside                  # visible without opening anything
    leaves_out = re.search(r"<details><summary>What this toy leaves out</summary>(.*?)</details>", ROOT, re.S).group(1)
    for said in ("Real guesses are not independent chances", "its own memory", "usually runs further below the chip", "the check stops being free",
                 "Llama 3.1 8B (big) and Llama 3.2 1B (draft) at 16 bits on one A100"):
        assert said in leaves_out, said
    assert re.findall(r'<li><\w+ data-kind="(\w+)">token</\w+><span>(.*?)</span></li>', ROOT) == [
        ("kept", "Guess kept"), ("rejected", "Guess rejected"), ("own", "The big model&rsquo;s own")]
    assert 'class="sdw-sr" id="sdw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".sdw-root .sdw-btns button { flex: 1 1 auto; white-space: nowrap; }" in HTML      # a button wraps to the next row, never its label
    assert "[data-kind=rejected]" in HTML and "line-through" in HTML                                          # a rejected guess is struck, not only orange


def test_widget_model_gives_the_random_numbers_rounds_and_sentences_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from speculative_sim.py: the lesson prints Python's numbers and the reader sees the widget's, so the
    random numbers, the rounds, every sentence and the two times must agree on every acceptance, guesses and draft setting the sliders offer."""
    cases = [{"pct": pct, "k": k, "d": d, "seed": 1 + i % 7, "view": settings_view(pct, k, d, 1 + i % 7)} for i, (pct, k, d) in enumerate(GRID)]
    seeds = (1, 2, 3, 12345, 2**32 - 1)
    streams = []
    for seed in seeds:
        draw = sim.random_stream(seed)
        streams.append({"seed": seed, "draws": [draw() for _ in range(50)]})
    expected = [[a, k, sim.expected_tokens(a, k)] for a in (0, 0.05, 0.5, 0.8, 0.95, 1) for k in (1, 4, 8)]
    result = run_node(tmp_path, PARITY, {"cases": cases, "streams": streams, "expected": expected}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] == 3 * len(cases) + len(streams) + len(expected)
    assert {c["view"]["verdict"] for c in cases} >= {"", "On average, about the same speed as plain decoding.", SLOWER}
    for c in cases:                                       # the verdict follows the averaged speed-up printed below it, not the one seeded run
        shown = float(re.search(r"([\d.]+) times as fast\.$", c["view"]["average"]).group(1))
        assert c["view"]["verdict"] == (SLOWER if shown < 1 else "On average, about the same speed as plain decoding." if shown == 1 else ""), c


def test_widget_page_shows_the_words_the_lesson_describes_and_replays_them(tmp_path):
    """Guards the part a reader sees: the opening state (80%, 4 guesses, 5 big passes against 17), the words in each lane spelling the reply, every kind
    of verdict, what is spoken on release and after New run, Replay frame by frame (the draft lane finishes first), a slider stopping it, reduced
    motion hiding it, and the theme message."""
    opening = settings_view(80, 4, sim.DEFAULT_DRAFT_STOP, sim.OPENING["seed"])
    assert sim.OPENING == {"percent": 80, "k": 4, "seed": 60} and opening["spec"].startswith("This reply: 5 big passes, 3.4 tokens each")
    assert len(opening["rounds"]) * opening["roundMs"] < len(sim.REPLY) * opening["plainMs"]
    verdicts = {}
    for pct, k, d in GRID:
        verdicts.setdefault(settings_view(pct, k, d, sim.OPENING["seed"])["verdict"], (pct, k, d))
    picks = [(pct, k, d) for pct in (0, 20, 50, 80, 100) for k in (1, 4, 8) for d in (0, 2, 5)] + list(verdicts.values())
    combos = [{"pct": pct, "k": k, "d": d, "view": settings_view(pct, k, d, sim.OPENING["seed"])} for pct, k, d in dict.fromkeys(picks)]
    new_runs = [{"pct": pct, "k": k, "d": 2, "view": settings_view(pct, k, 2, sim.OPENING["seed"] + 1)} for pct, k in ((80, 4), (20, 4), (100, 1))]
    ids = sorted(set(re.findall(r'id="(sdw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "reply": sim.REPLY, "decoys": sim.DECOYS, "defaultDraft": sim.DEFAULT_DRAFT_STOP, "combos": combos, "newRuns": new_runs,
               "opening": {"pct": 80, "k": 4, "view": opening}}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"combos": len(combos), "bad": []}
    assert len(verdicts) == 3