"""Guards for ../widget.html: its data block, its size, and that its JavaScript computes what gpu_split.py does."""
import json
import re
import subprocess

import gpu_split as g
from embed_widget import WIDGET, embed

HTML = WIDGET.read_text(encoding="utf-8")

# Reads the number code out of the widget between its marker comments, runs it, and compares with the
# values Python wrote into the fixture: the numbers, the usable memory and compute constants, the fits verdict,
# the bar widths and the chart rows the page draws.
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 DATA = JSON.parse(html.match(/<script type="application\/json" id="gpu-data">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the numbers'), html.indexOf('/* ---------- the page'));
const api = new Function('DATA', "'use strict';" + js + '; return {params, memory, fits, wordTime, usableBytes, computeTime, view, chart};')(DATA);
const close = (a, b) => Math.abs(a - b) <= 1e-9 * Math.abs(b);
const bad = [];
if (api.usableBytes() !== fx.usable) bad.push(['usableBytes', api.usableBytes(), fx.usable]);
fx.cases.forEach((c) => {
  const m = DATA.models.find((x) => x.key === c.key);
  const link = DATA.links.find((x) => x.key === c.link);
  if (api.params(m) !== c.params) bad.push(['params', c.key]);
  const mem = api.memory(m, c.n);
  for (const k of ['weights', 'cache', 'working', 'total']) if (mem[k] !== c.memory[k]) bad.push([c.key, c.n, k, mem[k], c.memory[k]]);
  if (api.fits(m, c.n) !== c.fits) bad.push([c.key, c.n, 'fits']);
  const t = api.wordTime(m, c.n, link, c.split, c.words);
  for (const [a, b] of [['read', 'read'], ['transfer', 'transfer'], ['total', 'total'], ['share', 'share'], ['wordsPerS', 'words_per_s']]) {
    if (!close(t[a], c.time[b])) bad.push([c.key, c.n, c.link, c.split, c.words, a, t[a], c.time[b]]);
  }
  if (!close(api.computeTime(m, c.n, c.words), c.compute)) bad.push([c.key, c.n, c.words, 'computeTime']);
  const v = api.view(m, c.n, link, c.split, c.words);
  if (v.ok !== c.fits) bad.push([c.key, c.n, c.link, c.split, c.words, 'view.ok', v.ok, c.fits]);
  if (!close(v.memBar, 100 * Math.min(1, c.memory.total / fx.usable))) bad.push([c.key, c.n, 'memBar', v.memBar, 100 * Math.min(1, c.memory.total / fx.usable)]);
  if (!close(v.waitBar, c.fits ? 100 * c.time.share : 0)) bad.push([c.key, c.n, c.link, c.split, c.words, 'waitBar', v.waitBar, c.fits ? 100 * c.time.share : 0]);
  // the chart for this model, link, split and words: one row per GPU count, bars against the slowest setting that fits
  const same = fx.cases.filter((x) => x.key === c.key && x.link === c.link && x.split === c.split && x.words === c.words);
  const top = Math.max(...same.filter((x) => x.fits).map((x) => x.time.total));
  const rows = api.chart(m, link, c.split, c.words);
  rows.forEach((r, i) => {
    const x = same[i];
    if (r.n !== x.n || r.ok !== x.fits) bad.push([c.key, c.link, c.split, c.words, 'chart row', r.n, r.ok, x && x.n, x && x.fits]);
    else if (x.fits && !(close(r.readBar, 100 * x.time.read / top) && close(r.waitBar, 100 * x.time.transfer / top))) bad.push([c.key, c.link, c.split, c.words, 'chart bars', r.n]);
  });
});
console.log(JSON.stringify({ cases: fx.cases.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def test_widget_data_block_matches_gpu_split():
    """Guards an edit to the constants or model files that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_every_element_the_script_looks_up_exists():
    """Guards a widget that never starts: a lookup of an id the page does not contain throws in the browser."""
    page = HTML[:HTML.index("<script")]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):]))
    assert lookups and all(f'id="{i}"' in page for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in page)
    assert 'id="gpu-data"' in HTML and 'id="gpu-root"' in page


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


def test_widget_numbers_match_the_python_code(tmp_path):
    """Guards the widget's JavaScript drifting from gpu_split.py: the usable memory and compute constants, memory, fit verdict,
    every time, and the bars and chart rows the page draws (they must follow the link and the split) must be the lesson's."""
    cases = [{"key": key, "n": n, "link": link, "split": split, "words": words, "params": g.param_count(g.load_config(key)),
              "memory": g.memory_per_gpu(key, n), "fits": g.fits(key, n), "time": g.word_time(key, n, link, split, words),
              "compute": g.compute_time(key, n, words)}
             for key, n, link, split, words in g.settings()]
    (tmp_path / "fixture.json").write_text(json.dumps({"usable": g.USABLE_BYTES, "cases": cases}), encoding="utf-8")
    (tmp_path / "parity.mjs").write_text(PARITY, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "parity.mjs"), str(tmp_path / "fixture.json"), str(WIDGET)],
                         capture_output=True, text=True, timeout=120)
    assert run.returncode == 0, run.stdout + run.stderr
    assert json.loads(run.stdout)["cases"] == len(g.MODELS) * len(g.COUNTS) * len(g.LINKS) * len(g.SPLITS) * len(g.WORDS) == 180
    assert any(c["fits"] for c in cases) and any(not c["fits"] for c in cases)