"""Guards for ../widget.html: its data block, its size, that its scheduler gives the same steps and words as batching_sim.py, and what a reader sees."""
import json
import math
import re
import subprocess
from functools import lru_cache

import batching_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="bw-root">'):HTML.index("<script")]
POLICIES = ("fixed", "continuous")
SETTINGS = [(rate, mix, slots) for mix in sim.MIXES for slots in sim.SLOT_CHOICES for rate in sim.RATES]
TICKS = [0, 1, 3, 250, 400, 500, 777, 999, 1000]

# Reads the scheduler out of the widget between its marker comments, runs it with Node, and compares every step, time, count and 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="bw-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 {makeRequests, simulate, report, snapshot, commas, fmtSeconds, timeAt, nowText, stateLine, summaryLine, stepSeconds};')(cfg);
const bad = [];
let checks = 0;
const same = (what, got, want) => { checks++; if (JSON.stringify(got) !== JSON.stringify(want)) bad.push([what, JSON.stringify(got).slice(0, 200), JSON.stringify(want).slice(0, 200)]); };
fx.cases.forEach((c) => {
  const what = c.rate + '/' + c.mix + '/' + c.slots + '/' + c.policy;
  const reqs = api.makeRequests(c.rate, c.mix);
  const run = api.simulate(reqs, c.slots, c.policy);
  same(what + ' arrivals', reqs.map((r) => r.arrival), c.arrivals);
  same(what + ' steps', run.steps.map((s) => [s.start, s.end, s.slots, s.joined, s.finished]), c.steps);
  same(what + ' first tokens', run.first, c.first);
  same(what + ' finish times', run.finished, c.finished);
  same(what + ' join steps', run.joinStep, c.joinStep);
  const rep = api.report(run, reqs, c.slots);
  same(what + ' report', [rep.seconds, rep.tokensPerSecond, rep.averageWait, rep.slowestWait, rep.slotsInUse], c.report);
  same(what + ' summary', api.summaryLine(rep), c.summary);
  c.snaps.forEach((s) => {
    const snap = api.snapshot(run, reqs, s.t);
    same(what + ' snapshot at ' + s.t, [snap.slots, snap.inUse, snap.waiting, snap.done], s.snap);
    same(what + ' state line at ' + s.t, api.stateLine(snap, c.slots, reqs.length), s.line);
  });
});
fx.words.forEach((w) => same('words ' + JSON.stringify(w.args), [api.commas(w.args[0]), api.fmtSeconds(w.args[1]), api.nowText(w.args[1], w.args[2]), api.timeAt(w.args[3], w.args[2])], w.want));
same('cost of one step', [api.stepSeconds(0, 0, 100), api.stepSeconds(1, 1000, 0), api.stepSeconds(64, 64000, 0)], fx.steps);
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 sliders, presses Play, changes the settings 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="bw-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, hidden: false, style: {}, attrs: {}, textContent: '' };
  e.setAttribute = (k, v) => { e.attrs[k] = v; };
  e.getAttribute = (k) => e.attrs[k];
  e.addEventListener = (type, f) => (e.listeners[type] = e.listeners[type] || []).push(f);
  e.appendChild = (c) => (e.children.push(c), c);
  e.replaceChildren = (...c) => { e.children = c; };
  return e;
}
function open(reduced) {
  const page = { root: el('html'), messages: [], frames: [], parent: { other: true } };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['bw-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f), matchMedia: () => ({ matches: !!reduced }), requestAnimationFrame: (f) => page.frames.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 = [];
let checks = 0;
const want = (what, got, expected) => { checks++; if (JSON.stringify(got) !== JSON.stringify(expected)) bad.push([what, JSON.stringify(got).slice(0, 160), JSON.stringify(expected).slice(0, 160)]); };
function look(page, v, what) {                       // compares the page with one expected view
  want(what + ': traffic words', page['bw-rate'].textContent, v.rate);
  want(what + ': time words', page['bw-now'].textContent, v.now);
  ['fixed', 'continuous'].forEach((name, k) => {
    const lane = v.lanes[k], cells = page['bw-' + name + '-cells'].children;
    want(what + ': ' + name + ' cells', cells.map((c) => [c.attrs['data-state'], c.children[0].style.width]), lane.cells.map((c) => [c.state, c.width]));
    want(what + ': ' + name + ' queue', page['bw-' + name + '-queue'].children.map((q) => q.attrs['data-on']), lane.squares);
    want(what + ': ' + name + ' more', page['bw-' + name + '-more'].textContent, lane.more);
    want(what + ': ' + name + ' state', page['bw-' + name + '-state'].textContent, lane.state);
    want(what + ': ' + name + ' summary', page['bw-' + name + '-summary'].textContent, lane.summary);
  });
}
const key = (rate, mix, slots, tick) => [rate, mix, slots, tick].join('|');
const spoken = (page) => page['bw-live'].textContent;

const first = open(false);
look(first, fx.views[key(8, 'mixed', 8, 400)], 'opening state');
want('opening controls', [first['bw-traffic'].value, first['bw-traffic'].max, first['bw-time'].value, first['bw-time'].max, first['bw-slots'].value, first['bw-mix'].value],
  [String(fx.rates.indexOf(8)), String(fx.rates.length - 1), '400', '1000', '8', 'mixed']);
want('option lists', [first['bw-slots'].children.map((o) => o.textContent), first['bw-mix'].children.map((o) => o.textContent)], [fx.slotChoices.map(String), fx.mixLabels]);
want('nothing spoken yet', spoken(first), '');
want('play is shown', first['bw-play'].hidden, false);

// every traffic stop, by dragging (input) and by releasing (change); the time position stays where it was
const traffic = open(false);
fx.rates.forEach((rate, i) => {
  traffic['bw-traffic'].value = String(i);
  const heard = spoken(traffic);
  fire(traffic['bw-traffic'], 'input');
  look(traffic, fx.views[key(rate, 'mixed', 8, 400)], 'traffic ' + rate);
  want('dragging stays silent at ' + rate, spoken(traffic), heard);
  want('slider reads out at ' + rate, traffic['bw-traffic'].attrs['aria-valuetext'], rate + ' new requests a second');
  fire(traffic['bw-traffic'], 'change');
  want('release speaks at ' + rate, spoken(traffic), fx.said.traffic[String(rate)]);
});

// the time slider at both ends and in between, and the words it speaks on release
const clock = open(false);
fx.ticks.forEach((tick) => {
  clock['bw-time'].value = String(tick);
  fire(clock['bw-time'], 'input');
  look(clock, fx.views[key(8, 'mixed', 8, tick)], 'tick ' + tick);
  fire(clock['bw-time'], 'change');
  want('time speaks at ' + tick, spoken(clock), fx.said.time[String(tick)]);
});

// the two settings under More settings
const more = open(false);
fx.slotChoices.forEach((n) => {
  more['bw-slots'].value = String(n);
  fire(more['bw-slots'], 'change');
  look(more, fx.views[key(8, 'mixed', n, 400)], 'slots ' + n);
  want('slots speak at ' + n, spoken(more), fx.said.slots[String(n)]);
});
more['bw-slots'].value = '8';
fx.mixes.forEach((mix) => {
  more['bw-mix'].value = mix;
  fire(more['bw-mix'], 'change');
  look(more, fx.views[key(8, mix, 8, 400)], 'mix ' + mix);
  want('mix speaks for ' + mix, spoken(more), fx.said.mix[mix]);
});

// Play: the first press plays the run from time 0; it moves the clock by the real time that passed, stops itself at the end, restarts from the start,
// goes on from where the clock is after Pause or a moved slider, and Pause and the slider stop it
const player = open(false);
fire(player['bw-play']);
want('play starts', [player['bw-play'].textContent, spoken(player), player.frames.length], ['Pause', 'Playing.', 1]);
let frames = player.frames.splice(0); frames[0](0);
look(player, fx.views[key(8, 'mixed', 8, 0)], 'first Play starts at time 0, not at the opening 40%');
frames = player.frames.splice(0); frames[0](1000);          // one second of real time is a tenth of the longer run
look(player, fx.views[key(8, 'mixed', 8, 100)], 'one second played');
want('time slider follows the clock', player['bw-time'].value, '100');
fire(player['bw-play']);
want('pause', [player['bw-play'].textContent, spoken(player)], ['Play', fx.said.time['100']]);
player.frames.splice(0);
fire(player['bw-play']);
frames = player.frames.splice(0); frames[0](5000);
look(player, fx.views[key(8, 'mixed', 8, 100)], 'Play after Pause goes on from the pause');
frames = player.frames.splice(0); frames[0](15000);
look(player, fx.views[key(8, 'mixed', 8, 1000)], 'end of the run');
want('stops itself at the end', [player['bw-play'].textContent, player.frames.length, spoken(player)], ['Play', 0, fx.said.finished]);
fire(player['bw-play']);
frames = player.frames.splice(0); frames[0](20000);
look(player, fx.views[key(8, 'mixed', 8, 0)], 'Play at the end starts over');
want('playing again', player['bw-play'].textContent, 'Pause');
player['bw-time'].value = '250'; fire(player['bw-time'], 'input');
want('moving the slider stops Play', [player['bw-play'].textContent, player['bw-time'].value], ['Play', '250']);
const dragged = open(false);                                  // a reader who put the clock back on the opening spot is still going on from there
dragged['bw-time'].value = '400'; fire(dragged['bw-time'], 'input');
fire(dragged['bw-play']);
frames = dragged.frames.splice(0); frames[0](0); frames = dragged.frames.splice(0); frames[0](1000);
look(dragged, fx.views[key(8, 'mixed', 8, 500)], 'Play after a moved slider goes on from the slider');

// reduced motion: no Play, nothing moves, and the page opens on the finished run
const calm = open(true);
want('reduced motion hides Play', calm['bw-play'].hidden, true);
look(calm, fx.views[key(8, 'mixed', 8, 1000)], 'reduced motion opens on the finished state');
want('reduced motion time position', calm['bw-time'].value, '1000');

// 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({ checks, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


@lru_cache(maxsize=None)
def runs(rate, mix, slots):
    requests = sim.make_requests(rate, mix)
    simulated = [sim.simulate(requests, slots, policy) for policy in POLICIES]
    return requests, simulated, [sim.report(run, requests, slots) for run in simulated]


def view(rate, mix, slots, tick):
    """What the page should show after it has settled on these settings and this time position."""
    requests, simulated, reports = runs(rate, mix, slots)
    end = max(rep["seconds"] for rep in reports)
    t = sim.time_at(tick, end)
    lanes = []
    for run, rep in zip(simulated, reports):
        snap = sim.snapshot(run, requests, t)
        waiting = snap["waiting"]
        lanes.append({
            "cells": [{"state": "use" if x >= 0 else "wait" if waiting else "empty", "width": f"{math.floor(x * 100)}%" if x >= 0 else "0%"} for x in snap["slots"]],
            "squares": ["1" if i < waiting else "0" for i in range(sim.QUEUE_CAP)],
            "more": f"+{waiting - sim.QUEUE_CAP}" if waiting > sim.QUEUE_CAP else "",
            "state": sim.state_line(snap, slots, len(requests)), "summary": sim.summary_line(rep),
        })
    return {"rate": sim.rate_text(rate), "now": sim.now_text(t, end), "lanes": lanes}


def said_after_a_run_setting(rate, mix, slots):
    _, _, reports = runs(rate, mix, slots)
    return f"{sim.rate_text(rate)} new requests a second. Fixed groups. {sim.summary_line(reports[0])} Continuous batching. {sim.summary_line(reports[1])}"


def said_at(tick, rate=8, mix="mixed", slots=8):
    v = view(rate, mix, slots, tick)
    return f"At {v['now']}. Fixed groups: {v['lanes'][0]['state']} Continuous batching: {v['lanes'][1]['state']}"


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")
    result = subprocess.run(["node", str(tmp_path / "check.mjs"), *args(tmp_path)], capture_output=True, text=True, timeout=300)
    assert result.returncode == 0, result.stdout + result.stderr
    return json.loads(result.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 a changed reply length must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"mixed": [', '"mixed": [999, ', 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 scheduler, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    for suffix in ("cells", "queue", "more", "state", "summary"):                    # looked up per lane as $('bw-' + name + '-suffix')
        assert f"$('bw-' + name + '-{suffix}')" in script or f"$('bw-' + lane.name + '-{suffix}')" in script, suffix
        lookups |= {f"bw-{lane}-{suffix}" for lane in POLICIES}
    missing = sorted(i for i in lookups if i != "bw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="bw-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 two sliders and Play stay in view, and the two selects wait 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)) == 3
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?<select id=\"bw-slots\">.*?<select id=\"bw-mix\">", 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)",
                 "--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", "#00855F", "#009E73", "#b36200", "#E69F00"}


def test_widget_scheduler_gives_the_steps_and_words_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from batching_sim.py: the lesson prints Python's numbers and the reader sees the widget's, so every step of every run
    (all 72 settings, both policies) must agree to the last bit, and so must the totals, the snapshots at nine moments and the lines of words."""
    cases = []
    for rate, mix, slots in SETTINGS:
        requests, simulated, reports = runs(rate, mix, slots)
        end = max(rep["seconds"] for rep in reports)
        for policy, run, rep in zip(POLICIES, simulated, reports):
            snaps = []
            for tick in TICKS:
                t = sim.time_at(tick, end)
                snap = sim.snapshot(run, requests, t)
                snaps.append({"t": t, "snap": [snap["slots"], snap["in_use"], snap["waiting"], snap["done"]], "line": sim.state_line(snap, slots, len(requests))})
            cases.append({"rate": rate, "mix": mix, "slots": slots, "policy": policy, "arrivals": [r.arrival for r in requests],
                          "steps": [[s.start, s.end, s.slots, s.joined, s.finished] for s in run.steps], "first": run.first_token, "finished": run.finished,
                          "joinStep": run.join_step, "snaps": snaps, "summary": sim.summary_line(rep),
                          "report": [rep["seconds"], rep["tokens_per_second"], rep["average_wait"], rep["slowest_wait"], rep["slots_in_use"]]})
    words = [{"args": [n, s, e, k], "want": [sim.commas(n), sim.fmt_seconds(s), sim.now_text(s, e), sim.time_at(k, e)]}
             for n, s, e, k in ((7, 0.0, 14.1, 0), (1234, 9.995, 103.9, 400), (999999, 10.0, 8.45, 1000), (0, 0.034, 0.5, 1), (12, 3.4349, 3.4349, 777))]
    steps = [sim.batch_step([], 100).seconds, sim.batch_step([1000]).seconds, sim.batch_step([1000] * 64).seconds]
    result = run_node(tmp_path, PARITY, {"cases": cases, "words": words, "steps": steps}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] >= len(cases) * (6 + len(TICKS) * 2 + 1)


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: the opening state (8 requests a second, 40% in), every traffic stop, the time slider, the two settings, Play and Pause, the
    finished state under reduced motion, what is spoken on release, the theme message, and the words that keep a best case from reading as a measurement."""
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert "A model at the chip&rsquo;s peak speed, not a measurement. The 48 requests are invented." in outside and sim.N_REQUESTS == 48
    assert "A new group starts only when every slot is empty." in outside and "A waiting request takes any empty slot at once." 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", "running out of memory for the caches", "Real servers are slower than these times", "Llama 3.1 8B with 16-bit weights on one A100"):
        assert said in re.sub(r"\s+", " ", leaves_out), said
    assert re.findall(r'<li><i data-state="(\w+)"></i><span>(.*?)</span></li>', ROOT) == [("use", "Reply under way"), ("empty", "Empty slot"), ("wait", "Empty while people wait")]
    assert 'class="bw-sr" id="bw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".bw-root button { font: inherit; font-weight: 600; font-size: 0.9rem; line-height: 1.2; min-height: 44px; min-width: 72px; padding: 2px 10px; white-space: nowrap;" in HTML
    assert ".bw-root [hidden] { display: none; }" in HTML                                              # Play really disappears under reduced motion
    views = {}
    mixes = list(sim.MIXES)
    for rate in sim.RATES:
        views[f"{rate}|mixed|8|400"] = view(rate, "mixed", 8, 400)
    for tick in TICKS + [100]:
        views[f"8|mixed|8|{tick}"] = view(8, "mixed", 8, tick)
    for slots in sim.SLOT_CHOICES:
        views[f"8|mixed|{slots}|400"] = view(8, "mixed", slots, 400)
    for mix in mixes:
        views[f"8|{mix}|8|400"] = view(8, mix, 8, 400)
    opening = views["8|mixed|8|400"]
    assert opening["lanes"][0]["state"] == "In use 2 of 8. 22 waiting. 15 of 48 done." and opening["lanes"][1]["state"] == "In use 7 of 8. 0 waiting. 32 of 48 done."
    assert opening["lanes"][0]["more"] == "+6" and opening["lanes"][1]["more"] == "" and opening["now"] == "5.63 s of 14.1 s"
    assert {c["state"] for c in opening["lanes"][0]["cells"]} == {"use", "wait"} and {c["state"] for c in opening["lanes"][1]["cells"]} == {"use", "empty"}
    said = {"traffic": {str(rate): said_after_a_run_setting(rate, "mixed", 8) for rate in sim.RATES},
            "time": {str(tick): said_at(tick) for tick in TICKS + [100]},
            "slots": {str(n): said_after_a_run_setting(8, "mixed", n) for n in sim.SLOT_CHOICES},
            "mix": {mix: said_after_a_run_setting(8, mix, 8) for mix in mixes},
            "finished": "Finished. " + said_at(1000)[said_at(1000).index("Fixed groups:"):]}
    ids = sorted(set(re.findall(r'id="(bw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "views": views, "said": said, "rates": sim.RATES, "slotChoices": sim.SLOT_CHOICES, "mixes": mixes, "ticks": TICKS,
               "mixLabels": [sim.mix_names()[m] for m in mixes]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result["bad"] == [] and result["checks"] > 200
