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

import embed_widget
import memory_calc as mc
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.
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="wmm-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, need, usable, fits, cards};')(DATA);
const bad = [];
fx.models.forEach((e) => { if (api.params(DATA.models.find((m) => m.key === e.key)) !== e.params) bad.push(['params', e.key]); });
fx.cases.forEach((c) => {
  const m = DATA.models.find((x) => x.key === c.key);
  const p = api.need(m, c.bits, c.tokens, c.chats);
  for (const [a, b] of [['weights', 'weights'], ['kv', 'kv_cache'], ['act', 'activations'], ['overhead', 'overhead'], ['total', 'total']]) {
    if (p[a] !== c.parts[b]) bad.push([c.key, c.bits, c.tokens, c.chats, a, p[a], c.parts[b]]);
  }
  c.gpus.forEach((g) => {
    if (api.fits(p.total, g.gib) !== g.fits || api.cards(p.total, g.gib) !== g.cards) bad.push([c.key, c.bits, c.tokens, c.chats, 'gpu', g]);
  });
});
fx.edges.forEach((e) => {
  if (api.fits(e.total, e.gib) !== e.fits || api.cards(e.total, e.gib) !== e.cards) bad.push(['edge', e]);
});
console.log(JSON.stringify({ cases: fx.cases.length, edges: fx.edges.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def test_widget_data_block_matches_memory_calc():
    """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="wmm-data"' in HTML and 'id="wmm-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 memory_calc.py: its sizes and its fit verdicts must be the lesson's."""
    models = [{"key": k, "params": mc.param_count(mc.load_config(k))} for k in mc.MODELS]
    cases = []
    for key in mc.MODELS:
        for bits in mc.BITS:
            for tokens in mc.CONTEXTS:
                for chats in mc.CHATS:
                    parts = mc.memory_need(key, bits, tokens, chats)
                    cases.append({"key": key, "bits": bits, "tokens": tokens, "chats": chats, "parts": parts,
                                  "gpus": [{"gib": g, "fits": mc.fits(parts["total"], g), "cards": mc.cards_needed(parts["total"], g)}
                                           for g, _ in mc.GPUS]})
    # Totals one byte either side of each card's usable limit, and of two cards' worth: 24 x 2**30 x 0.92 is not a whole number of bytes.
    edges = [{"gib": g, "total": t, "fits": mc.fits(t, g), "cards": mc.cards_needed(t, g)}
             for g, _ in mc.GPUS for u in [mc.usable_bytes(g)] for t in (u - 1, u, u + 1, 2 * u, 2 * u + 1)]
    (tmp_path / "fixture.json").write_text(json.dumps({"models": models, "cases": cases, "edges": edges}), 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(mc.MODELS) * len(mc.BITS) * len(mc.CONTEXTS) * len(mc.CHATS) == 756
    assert json.loads(run.stdout)["edges"] == 25
    assert any(not g["fits"] for c in cases for g in c["gpus"]) and any(g["fits"] for c in cases for g in c["gpus"])
    assert [e["fits"] for e in edges[:5]] == [True, True, False, False, False]   # the edge itself fits, one byte more does not