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

import numpy as np

import quantize as qz
from embed_widget import WIDGET, embed

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

# Reads the number code out of the widget between its marker comments, runs it, and compares with the
# values Python wrote into the fixture. Rounding must agree to the last bit; only sums may differ by rounding noise.
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="qnt-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 {rne, f16, quantize, sampleError, weightBytes, bytesPerToken, ceiling};')(DATA);
const bad = [];
const close = (a, b) => Math.abs(a - b) <= 1e-12 * Math.max(1, Math.abs(b));
fx.f16.forEach(([x, want]) => { if (api.f16(x) !== want) bad.push(['f16', x, api.f16(x), want]); });
fx.rne.forEach(([x, want]) => { if (api.rne(x) !== want) bad.push(['rne', x, api.rne(x), want]); });
fx.quant.forEach((c) => {
  const q = api.quantize(c.w, c.bits, c.group);
  if (JSON.stringify(q.codes) !== JSON.stringify(c.codes) || JSON.stringify(q.scales) !== JSON.stringify(c.scales)) bad.push(['quantize', c.bits, c.group, c.w, q.codes, c.codes]);
  if (!q.back.every((v, i) => close(v, c.back[i]))) bad.push(['back', c.bits, c.group]);
});
fx.sample.forEach((c) => { if (!close(api.sampleError(c.bits, c.group, c.outlier), c.error)) bad.push(['sampleError', c]); });
fx.model.forEach((c) => {
  if (api.weightBytes(c.bits, c.group) !== c.weights || api.bytesPerToken(c.bits, c.group) !== c.read) bad.push(['bytes', c]);
  c.gpus.forEach((g) => { if (api.ceiling(c.bits, c.group, g.bandwidth) !== g.tps) bad.push(['tps', c.bits, c.group, g]); });
});
console.log(JSON.stringify({ f16: fx.f16.length, quant: fx.quant.length, sample: fx.sample.length, model: fx.model.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def test_widget_data_block_matches_quantize():
    """Guards an edit to the constants or sample weights 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="qnt-data"' in HTML and 'id="qnt-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 quantize.py: rounding (ties to even), 16-bit scales, errors, sizes and speeds."""
    rng = np.random.default_rng(5)
    f16 = [[float(x), float(np.float16(x))] for x in np.concatenate([10 ** rng.uniform(-9, 3, 400), [0.0, 1.0, 2.0**-14, 2.0**-24, 1.0009765625, 65504.0]])]
    rne = [[x, float(np.rint(x))] for x in [-2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, -0.4, 0.51, 7.0] + list(rng.uniform(-130, 130, 200))]
    quant = []
    vectors = [qz.sample_weights(), qz.sample_with_outlier(), [3.0, 2.5, 0.5, -0.5, -1.5, -2.5, 1.0, 0.0]]
    vectors += [list(np.round(rng.normal(0, s, 16), 4)) for s in (0.001, 0.1, 5.0) for _ in range(20)]
    for w in vectors:
        for bits in (8, 4, 3):
            for group in (None, 4, 8):
                codes, scales = qz.quantize(w, bits, group)
                quant.append({"w": [float(x) for x in w], "bits": bits, "group": group, "codes": codes.ravel().tolist(),
                              "scales": scales.astype(np.float64).ravel().tolist(), "back": qz.dequantize(codes, scales, group).ravel().tolist()})
    sample = [{"bits": b, "group": g, "outlier": o, "error": qz.sample_error(b, g, o)}
              for b in qz.BITS for g in (None, qz.SAMPLE_GROUP) for o in (False, True)]
    model = [{"bits": b, "group": g, "weights": qz.weight_bytes(CFG, b, g), "read": qz.bytes_per_token(CFG, b, g),
              "gpus": [{"bandwidth": bw, "tps": qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b, g), bw)} for bw in qz.BANDWIDTHS.values()]}
             for b in qz.BITS for g in (None, qz.GROUP)]
    (tmp_path / "fixture.json").write_text(json.dumps({"f16": f16, "rne": rne, "quant": quant, "sample": sample, "model": model}), 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)["quant"] == len(vectors) * 9 == 567
    assert any(c["codes"].count(0) > 10 for c in quant) and any(max(map(abs, c["codes"])) == 127 for c in quant)


# Runs the widget's page code under Node against a minimal stand-in for the document, reads the three readouts, then clicks
# the outlier checkbox and picks "One for all" as the lesson tells the reader to.
RENDER = r"""
import fs from 'node:fs';
const html = fs.readFileSync(process.argv[2], 'utf8');
const json = html.match(/<script type="application\/json" id="qnt-data">([\s\S]*?)<\/script>/)[1];
const code = html.slice(html.indexOf('<script>', html.indexOf('</script>')) + 8, html.lastIndexOf('</script>'));
const els = {}, handlers = {};
const el = (id) => els[id] ||= { textContent: id === 'qnt-data' ? json : '', innerHTML: '', style: {}, checked: false,
  addEventListener(type, fn) { (handlers[id] ||= {})[type] = fn; }, getBoundingClientRect: () => ({ height: 0 }), querySelectorAll: () => [] };
const win = { addEventListener() {} };
new Function('document', 'window', 'parent', code)({ getElementById: el }, win, win);
const read = () => ({ size: els['qnt-size'].textContent, speed: els['qnt-speed'].textContent, error: els['qnt-err'].textContent });
const out = { open: read() };
handlers['qnt-big'].change({ target: { checked: true } });
out.outlier = read();
handlers['qnt-mode'].change({ target: { value: '0' } });
out.oneForAll = read();
console.log(JSON.stringify(out));
"""


def test_widget_readouts_follow_the_lesson(tmp_path):
    """Guards the page code the parity test cannot reach: size and speed always use groups of 128 (picking "One for all" changed
    them before), GiB is 2**30 bytes, and the opening state, the big-weight box and the scale choice reach the readouts."""
    (tmp_path / "render.mjs").write_text(RENDER, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "render.mjs"), str(WIDGET)], capture_output=True, text=True, timeout=60)
    assert run.returncode == 0, run.stdout + run.stderr
    out = json.loads(run.stdout)
    assert out["open"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "4.7%"}
    assert out["outlier"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "19.1%"}
    assert out["oneForAll"] == {"size": "5.31 GiB", "speed": "721 tokens/s", "error": "34.7%"}