"""Guards for ../widget.html: its data block, its size, that its formulas give the same times and words as prefill_sim.py, and what a reader sees."""
import json
import re
import subprocess

import pytest

import prefill_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="pfw-root">'):HTML.index("<script")]
STOPS = sim.prompt_lengths()
EXTRA_LENGTHS = [1, 2, 153, 154, 155, 5000, 6126, 100_000]       # beyond the slider, where a formula slip would show

# Reads the model out of the widget between its marker comments, runs it with Node, and compares every time and every line of words
# with the ones Python wrote into the fixture, bit for bit (the seconds are compared as exact numbers).
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="pfw-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 {pass: pass, prefill, decodeStep, view, commas, fmtMs};')(cfg);
const bad = [];
let checks = 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
fx.cases.forEach((c) => {
  const p = api.prefill(c.n, c.extras), d = api.decodeStep(c.n, c.extras);
  checks += 3;
  if (p.seconds !== c.prefill.seconds || p.bound !== c.prefill.bound || p.busy !== c.prefill.busy) bad.push(['prefill', c.n, c.extras, p, c.prefill]);
  if (d.seconds !== c.decode.seconds || d.bound !== c.decode.bound || d.busy !== c.decode.busy) bad.push(['decode', c.n, c.extras, d, c.decode]);
  if (!same(api.view(c.n, c.extras), c.view)) bad.push(['view', c.n, c.extras, api.view(c.n, c.extras), c.view]);
});
checks++;                                          // a tie at the ridge counts as arithmetic, as in prefill_sim.py
if (api.pass(cfg.peak, cfg.bandwidth).bound !== 'arithmetic') bad.push(['tie', api.pass(cfg.peak, cfg.bandwidth)]);
fx.commas.forEach(([n, text]) => { checks++; if (api.commas(n) !== text) bad.push(['commas', n, api.commas(n), text]); });
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, moves the slider, presses the buttons and the checkbox,
# and reads back what a reader sees.
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="pfw-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, value: '', checked: false, style: {}, attrs: {}, textContent: '' };
  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() {
  const page = { root: el('html'), messages: [], parent: { other: true } };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['pfw-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f) };
  new Function('document', 'window', 'parent', code)(doc, win, page.parent);
  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]); };
function look(page, v, what, n) {                    // compares the page with one expected view
  want(what + ': length', page['pfw-n'].textContent, n);
  ['ttft', 'tpot', 'why', 'busy'].forEach((k) => want(what + ': ' + k, page['pfw-' + k].textContent, v[k]));
  const blocks = page['pfw-bar'].children;
  want(what + ': block count', blocks.length, fx.drawnSteps + 1);
  want(what + ': prefill width', blocks[0].style.width, v.prefillWidth);
  want(what + ': prefill colour', blocks[0].attrs['data-bound'], v.bound);
  want(what + ': first word mark', page['pfw-mark'].style.left, v.prefillWidth);
  want(what + ': step widths', blocks.slice(1).map((b) => b.style.width), blocks.slice(1).map(() => v.stepWidth));
  want(what + ': step colours', blocks.slice(1).map((b) => b.attrs['data-bound']), blocks.slice(1).map(() => 'memory'));
}

const first = open();
look(first, fx.opening, 'opening state', fx.openingText);
want('slider range', [first['pfw-len'].max, first['pfw-len'].value], [String(fx.stops.length - 1), String(fx.stops.indexOf(fx.openingTokens))]);
want('axis labels', [first['pfw-mid'].textContent, first['pfw-end'].textContent], fx.axis);
want('nothing spoken yet', first['pfw-live'].textContent, '');
const buttons = first['pfw-btns'].children;
want('preset labels', buttons.map((b) => b.textContent), fx.presetLabels);
want('preset pressed at the opening', buttons.map((b) => b.attrs['aria-pressed']), fx.presets.map((p) => String(p.tokens === fx.openingTokens)));

// every slider stop, with the checkbox off and on, by dragging (input event) and by releasing (change event)
[false, true].forEach((extras) => {
  const page = open();
  page['pfw-extra'].checked = extras;
  fx.stops.forEach((n, i) => {
    page['pfw-len'].value = String(i);
    page['pfw-live'].textContent = '';
    fire(page['pfw-len'], 'input');
    const v = fx.views[String(extras)][String(n)];
    look(page, v, 'stop ' + n + ' extras ' + extras, fx.texts[String(n)]);
    want('dragging stays silent at ' + n, page['pfw-live'].textContent, '');
    want('slider reads out', page['pfw-len'].attrs['aria-valuetext'], fx.texts[String(n)] + ' tokens');
    fire(page['pfw-len'], 'change');
    want('release speaks at ' + n, page['pfw-live'].textContent, fx.texts[String(n)] + ' tokens. ' + v.ttft + '. ' + v.tpot);
  });
});

// the three presets move the slider and mark themselves pressed
const press = open();
fx.presets.forEach((p, k) => {
  fire(press['pfw-btns'].children[k]);
  const n = p.tokens;
  look(press, fx.views.false[String(n)], 'preset ' + p.name, fx.texts[String(n)]);
  want('preset ' + p.name + ' slider', press['pfw-len'].value, String(fx.stops.indexOf(n)));
  want('preset ' + p.name + ' pressed', press['pfw-btns'].children.map((b) => b.attrs['aria-pressed']), fx.presets.map((q, j) => String(j === k)));
  want('preset ' + p.name + ' spoken', press['pfw-live'].textContent, fx.texts[String(n)] + ' tokens. ' + fx.views.false[String(n)].ttft + '. ' + fx.views.false[String(n)].tpot);
});
// ticking the box recomputes the current length and speaks
press['pfw-extra'].checked = true; fire(press['pfw-extra'], 'change');
const longest = String(fx.presets[2].tokens), ticked = fx.views.true[longest];
look(press, ticked, 'box ticked at the longest prompt', fx.texts[longest]);
want('box ticked speaks', press['pfw-live'].textContent, fx.texts[longest] + ' tokens. ' + ticked.ttft + '. ' + ticked.tpot);

// the host page theme: only the parent window may set it, and only to light or dark
const themed = open();
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({ stops: fx.stops.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def pass_record(p):
    return {"seconds": p.seconds, "bound": p.bound, "busy": p.busy}


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 slider stop must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"axisMs": 600', '"axisMs": 500', 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 formulas, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script)) | {f"pfw-{k}" for k in ("ttft", "tpot", "why", "busy")}
    assert "$('pfw-' + k)" in script                      # the four text lines are looked up by suffix; they are added above
    missing = sorted(i for i in lookups if i != "pfw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="pfw-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: the slider and the preset buttons (built by the script) count, and the checkbox stays 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)) + len(sim.PRESETS) <= 5
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?type=\"checkbox\"", ROOT, re.S)


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_more_settings_note_quotes_what_the_code_computes():
    """Guards a stale sentence: the details text says the extras add about 13% to the first word and 6.5% to each later token at 8,000 tokens."""
    n = sim.MAX_PROMPT
    first = 100 * (sim.prefill(n, True).seconds / sim.prefill(n).seconds - 1)
    later = 100 * (sim.decode_step(n, True).seconds / sim.decode_step(n).seconds - 1)
    assert f"{first:.0f}%" == "13%" and f"{later:.1f}%" == "6.5%"
    assert "about 13% later and each later token about 6.5% later" in HTML and "At 8,000 tokens" in HTML


def test_widget_formulas_give_the_times_and_words_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from prefill_sim.py: the lesson prints Python's numbers and the reader sees the
    widget's, so the two must agree to the last bit on every slider stop, with the extras on and off, and on lengths beyond the slider."""
    cases = []
    for extras in (False, True):
        for n in STOPS + EXTRA_LENGTHS:
            cases.append({"n": n, "extras": extras, "prefill": pass_record(sim.prefill(n, extras)), "decode": pass_record(sim.decode_step(n, extras)),
                          "view": sim.view(n, extras)})
    commas = [[n, sim.commas(n)] for n in (0, 7, 999, 1000, 8000, 12345, 1234567)]
    result = run_node(tmp_path, PARITY, {"cases": cases, "commas": commas}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] == 3 * len(cases) + len(commas) + 1
    assert {c["view"]["bound"] for c in cases} == {"memory", "arithmetic"}


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: the opening state (8,000 tokens, 412 ms, a wide green block), every slider stop, the
    presets, the checkbox, what is spoken on release, the theme message, and the words that keep a best case from reading as a
    measurement (the visible caveat line, the closed panel, the legend named by limit, the live region)."""
    views = {str(extras).lower(): {str(n): sim.view(n, extras) for n in STOPS} for extras in (False, True)}
    opening = views["false"]["8000"]
    assert opening["ttft"] == "Best-case time to the first token: 412 ms" and opening["bound"] == "arithmetic"
    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
    assert "Left of the line: prefill, the whole prompt in one pass. Right of it: decode, one token per step." in outside
    leaves_out = re.search(r"<details><summary>What this toy leaves out</summary>(.*?)</details>", ROOT, re.S).group(1)
    for said in ("Network time, waiting in line behind other users", "Real servers are slower than these times", "Llama 3.1 8B with 16-bit weights on one A100"):
        assert said in leaves_out, said
    assert re.findall(r'<li><i data-bound="(\w+)"></i><span>(.*?)</span></li>', ROOT) == [("memory", "Limited by memory"), ("arithmetic", "Limited by arithmetic")]
    assert 'class="pfw-sr" id="pfw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".pfw-root .pfw-btns button { flex: 1 1 auto; white-space: nowrap; }" in HTML      # a button wraps to the next row, never its label
    assert views["false"]["20"]["bound"] == "memory" and views["false"]["20"]["prefillWidth"] == views["false"]["20"]["stepWidth"]
    ids = sorted(set(re.findall(r'id="(pfw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "stops": STOPS, "views": views, "texts": {str(n): sim.commas(n) for n in STOPS}, "opening": opening,
               "openingText": "8,000", "openingTokens": sim.OPENING_TOKENS, "drawnSteps": sim.DRAWN_STEPS, "presets": sim.PRESETS,
               "presetLabels": [f"{p['name']} ({sim.commas(p['tokens'])})" for p in sim.PRESETS], "axis": [f"{sim.AXIS_MS // 2} ms", f"{sim.AXIS_MS} ms"]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"stops": len(STOPS), "bad": []}