"""Guards for ../widget.html: its data block, its size, that its readers see the same opcodes and header as modelfile.py, and what a reader sees."""
import json
import pickle
import re
import struct
import subprocess

import pytest

import harmless
import modelfile as mf
from embed_widget import CONFIG_BLOCK, WIDGET, config, embed

HTML = WIDGET.read_text(encoding="utf-8")
ROOT = HTML[HTML.index('<div class="mf-root">'):HTML.index("<script")]
OUTSIDE = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)

# Reads the byte readers out of the widget between its marker comments, runs them with Node on files that Python built, and compares every
# opcode (position, end, name, argument), the step where the restricted loader stops, and every safetensors field, with what Python found.
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 js = html.slice(html.indexOf('/* ---------- readers ('), html.indexOf('/* ---------- the words'));
const api = new Function("'use strict';" + js + '; return {toBytes, num, readOps, firstRefused, readSafe};')();
const bad = [];
let checks = 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
fx.pickles.forEach((p) => {
  const ops = api.readOps(api.toBytes(p.hex));
  checks += 2;
  if (!same(ops.map((o) => [o.pos, o.end, o.name, o.arg]), p.ops)) bad.push(['ops', p.name, ops.slice(0, 3)]);
  if (api.firstRefused(ops, fx.allowed) !== p.first) bad.push(['first refused', p.name, api.firstRefused(ops, fx.allowed), p.first]);
});
fx.safe.forEach((s) => {
  const got = api.readSafe(api.toBytes(s.hex));
  checks++;
  if (!same(got, s.read)) bad.push(['safetensors', s.name, got, s.read]);
});
[[2, '2.0'], [1.5, '1.5'], [-0.25, '-0.25'], [0.125, '0.125'], [-1, '-1.0']].forEach(([x, t]) => { checks++; if (api.num(x) !== t) bad.push(['num', x, api.num(x)]); });
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="mf-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

let active = null;
function el(tag) {
  const e = { tag, children: [], listeners: {}, value: '', checked: false, disabled: false, style: {}, attrs: {}, focus: () => { active = e; } };
  let text = '';                                     // as in a browser, setting the text of an element drops its children
  Object.defineProperty(e, 'textContent', { get: () => text, set: (v) => { text = v; e.children.length = 0; } });
  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['mf-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root, get activeElement() { return active; } };
  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]); };
const states = (page, id) => page[id].children.map((c) => c.attrs['data-s']);
const lit = (pre) => pre.children.filter((c) => c.attrs['class'] === 'mf-on').map((c) => c.textContent);
const step = (page, n) => { page['mf-step'].value = String(n); fire(page['mf-step'], 'input'); };

const first = open();
want('opening: headline', first['mf-head'].textContent, fx.text.opening);
want('opening: slider', [first['mf-step'].value, first['mf-step'].max, first['mf-n'].textContent, first['mf-max'].textContent], ['28', '28', '28', '28']);
want('opening: sizes', [first['mf-psize'].textContent, first['mf-ssize'].textContent], [String(fx.pickleBytes), String(fx.safeBytes)]);
want('opening: counts', [first['mf-pcount'].textContent, first['mf-scount'].textContent], [fx.text.one, fx.text.none]);
want('opening: pickle strip', states(first, 'mf-pstrip'), Array(27).fill('done').concat(['now']));
want('opening: safetensors strip', states(first, 'mf-sstrip'), ['done', 'done', 'now']);
want('opening: kinds', first['mf-pstrip'].children.map((c) => c.attrs['data-k']).map((k, i) => (k === 'build' ? '' : i + k)).filter(Boolean), ['18import', '24call']);
want('opening: labels over the strip', first['mf-marks'].children.map((c) => c.textContent), ['import', 'call']);
want('opening: safetensors line', first['mf-sline'].textContent, fx.text.safeDone);
want('opening: next is off, nothing spoken', [first['mf-next'].disabled, first['mf-live'].textContent], [true, '']);

const page = open();
step(page, 19);
want('step 19: line', page['mf-pline'].textContent, fx.text.importLine);
want('step 19: headline', page['mf-head'].textContent, fx.text.after19);
want('step 19: strip', states(page, 'mf-pstrip').indexOf('now'), 18);
want('step 19: counts', page['mf-pcount'].textContent, fx.text.none);
want('step 19: bytes outlined', lit(page['mf-phex']), [fx.hex.importStep]);
want('dragging stays silent', page['mf-live'].textContent, '');
step(page, 25);
want('step 25: line', page['mf-pline'].textContent, fx.text.callLine);
want('step 25: counts', page['mf-pcount'].textContent, fx.text.one);
want('step 25: bytes outlined', lit(page['mf-phex']), [fx.hex.callStep]);
fire(page['mf-step'], 'change');
want('release speaks', page['mf-live'].textContent, page['mf-head'].textContent + ' ' + page['mf-pline'].textContent);
step(page, 2);
want('step 2: safetensors bytes', [page['mf-sline'].textContent, lit(page['mf-shex'])], [fx.text.safeStep2, [fx.hex.header]]);
want('step 2: safetensors strip', states(page, 'mf-sstrip'), ['done', 'now', 'todo']);

// every slider stop, with the restricted loader off and on: calls count REDUCE steps reached, and the stop is at the first refused import
[false, true].forEach((restricted) => {
  const p = open();
  p['mf-gate'].checked = restricted;
  for (let n = 0; n <= 28; n++) {
    step(p, n);
    const stopped = restricted && n > fx.first;
    const calls = stopped ? 0 : fx.calls[n];
    want('stop ' + n + ' restricted ' + restricted + ': counts', p['mf-pcount'].textContent, 'Functions called: ' + calls + ' · Notes recorded: ' + calls);
    want('stop ' + n + ': stopped words', p['mf-head'].textContent.startsWith('The restricted loader stopped'), stopped);
    want('stop ' + n + ': safetensors calls', p['mf-scount'].textContent, fx.text.none);
    want('stop ' + n + ': next', p['mf-next'].disabled, n >= 28 || stopped);
    if (stopped) {
      const s = states(p, 'mf-pstrip');
      want('stopped strip ' + n, [s[fx.first], s.slice(fx.first + 1).every((x) => x === 'skip'), s.slice(0, fx.first).every((x) => x === 'done')], ['stop', true, true]);
    }
  }
});

// the restricted loader at the end of the file says where it stopped and that nothing ran; unticking restores the opening
const gate = open();
gate['mf-gate'].checked = true; fire(gate['mf-gate'], 'change');
want('restricted: headline', gate['mf-head'].textContent, fx.text.stopped);
want('restricted: line', gate['mf-pline'].textContent, fx.text.stopLine);
want('restricted: speaks', gate['mf-live'].textContent, fx.text.stopped + ' ' + fx.text.stopLine);
want('restricted: safetensors unchanged', gate['mf-sline'].textContent, fx.text.safeDone);
gate['mf-gate'].checked = false; fire(gate['mf-gate'], 'change');
want('unticked: headline', gate['mf-head'].textContent, fx.text.opening);

// the restricted loader never shows a step past its stop: the counter, the spoken value and the slider say where it is, and Next is off
const stop = open();
stop['mf-gate'].checked = true; step(stop, 25);
want('restricted at 25: counter, spoken value, slider, next', [stop['mf-n'].textContent, stop['mf-step'].attrs['aria-valuetext'], stop['mf-step'].value, stop['mf-next'].disabled],
     [String(fx.first + 1), 'step ' + (fx.first + 1) + ' of 28', String(fx.first + 1), true]);

// the bytes outlined in the safetensors dump at its three reading steps: the length, the header, the numbers
[[1, 8], [2, 68], [3, 24]].forEach(([n, bytes]) => {
  step(page, n);
  want('safetensors step ' + n + ': bytes outlined', lit(page['mf-shex']).map((t) => t.split(' ').length), [bytes]);
});

// Next step and Start over
const btn = open();
fire(btn['mf-reset']);
want('start over', [btn['mf-step'].value, btn['mf-pline'].textContent, btn['mf-sline'].textContent, btn['mf-next'].disabled], ['0', fx.text.notStarted, 'Not started.', false]);
btn['mf-live'].textContent = '';                  // Start over spoke; only the click on Next may fill it now
fire(btn['mf-next']);
want('next advances by one step', [btn['mf-step'].value, btn['mf-n'].textContent, btn['mf-live'].textContent !== ''], ['1', '1', true]);
btn['mf-next'].focus();
for (let i = 0; i < 40; i++) fire(btn['mf-next']);
want('next stops at the end', [btn['mf-step'].value, btn['mf-next'].disabled, btn['mf-head'].textContent, active === btn['mf-step']], ['28', true, fx.text.opening, true]);

// 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' });

console.log(JSON.stringify({ bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""

PICKLES = {
    "demo": mf.demo_pickle(),
    "plain": mf.plain_pickle(),
    "range": pickle.dumps({"r": range(3), "s": {1, 2}}, protocol=3),
    "many kinds": pickle.dumps({"a": (1, None, True, False, 70000, -5, 300), "b": [1, 2, 3], "c": (1,), "d": (1, 2), "e": (), "f": 0.5, "g": "xé"}, protocol=3),
    "two imports": pickle.dumps([mf.Opens(), range(2)], protocol=3),
}


def hexed(data, begin, end):
    return " ".join(f"{b:02x}" for b in data[begin:end])


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 arg_text(arg):
    return "" if arg is None else repr(arg) if isinstance(arg, float) else str(arg)


def pickle_fixture(name, data):
    ops = mf.opcodes(data)
    ends = [o[0] for o in ops[1:]] + [len(data)]
    first = mf.first_refused(data)
    return {"name": name, "hex": data.hex(), "ops": [[pos, end, op, arg_text(arg)] for (pos, op, arg), end in zip(ops, ends)], "first": -1 if first is None else first}


def safe_fixture(name, data):
    tensors, _ = mf.read_safetensors(data)
    (n,) = struct.unpack_from("<Q", data, 0)
    header = json.loads(data[8:8 + n])
    read = {k: {"dtype": d, "shape": list(s), "offsets": header[k]["data_offsets"], "values": list(v)} for k, (d, s, v) in tensors.items()}
    return {"name": name, "hex": data.hex(), "read": {"headerBytes": n, "tensors": read}}


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 note or file must be put back."""
    drifted = HTML.replace(mf.NOTE, "something else", 1).replace('"builtins.range"', '"builtins.set"', 1)
    assert drifted != HTML and embed(drifted) == HTML
    assert json.loads(CONFIG_BLOCK.search(HTML).group(2)) == json.loads(json.dumps(config()))
    assert bytes.fromhex(config()["pickle"]) == mf.demo_pickle() and bytes.fromhex(config()["safetensors"]) == mf.demo_safetensors()


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: the parity run covers only the readers, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script)) | {f"mf-{k}" for k in ("head", "pline", "pcount", "sline", "scount")}
    assert "$('mf-' + k)" in script                        # the five text lines are looked up by suffix; they are added above
    missing = sorted(i for i in lookups if i != "mf-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="mf-config"' in HTML


def test_widget_is_small_asks_for_nothing_online_and_runs_no_text():
    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(", "new Function", "innerHTML", "<form", 'type="submit"'):
        assert banned not in page, banned
    assert " src=" not in page and "@import" not in page


def test_widget_shows_at_most_five_controls_and_hides_the_rest():
    """Guards the first screen at 390 px: the slider, the two buttons and the checkbox are all there is; the byte dumps sit in a closed details element."""
    assert len(re.findall(r"<(?:select|button|input|textarea)\b", OUTSIDE)) <= 5
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?id=\"mf-phex\".*?id=\"mf-shex\"", 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
    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_widget_says_what_it_does_and_does_not_do():
    """Guards the sentence the lesson promises: the page shows what Python's loader would do and runs nothing itself, visible without opening anything."""
    assert "The widget shows what Python&rsquo;s loader would do with the bytes; it only reads them and writes words, and runs nothing itself." in OUTSIDE
    leaves_out = re.search(r"<details><summary>What this toy leaves out</summary>(.*?)</details>", ROOT, re.S).group(1)
    for said in ("millions or billions of numbers", "you can add types to its allowed list", "could name any function Python can import", "no holes"):
        assert said in leaves_out, said
    assert ".mf-root .mf-btns button { flex: 1 1 auto; white-space: nowrap; }" in HTML      # a button wraps to the next row, never its label
    assert 'class="mf-sr" id="mf-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT


def test_widget_readers_see_what_python_sees(tmp_path):
    """Guards the widget's JavaScript drifting from modelfile.py: the lesson prints Python's numbers and the reader sees the widget's, so the opcode
    list, the step where the restricted loader stops and the safetensors header must agree on the demo files and on other pickles."""
    pickles = [pickle_fixture(name, data) for name, data in PICKLES.items()]
    safe = [safe_fixture("demo", mf.demo_safetensors()),
            safe_fixture("two tensors", mf.write_safetensors({"a": ("F32", (3,), (0.5, 1.0, -2.0)), "b": ("I32", (2,), (7, -9))}, {"format": "pt"}))]
    assert pickles[0]["first"] == 18 and pickles[1]["first"] == -1 and pickles[2]["first"] == -1 and pickles[4]["first"] > 0
    allowed = config()["allowed"]
    result = run_node(tmp_path, PARITY, {"pickles": pickles, "safe": safe, "allowed": allowed}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] == 2 * len(pickles) + len(safe) + 5


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: the opening state (one function called by the pickle, none by safetensors), the import at step 19 and the call at
    step 25, every slider stop with the restricted loader off and on, the stop cell, the buttons, what is spoken on release, and the theme message."""
    ops = mf.opcodes(mf.demo_pickle())
    calls = [sum(1 for o in ops[:n] if o[1] == "REDUCE") for n in range(len(ops) + 1)]
    n = mf.numbers()
    fx = {"ids": sorted(set(re.findall(r'id="(mf-[\w-]+)"', HTML))), "pickleBytes": n["pickle_bytes"], "safeBytes": n["safetensors_bytes"], "calls": calls,
          "first": mf.first_refused(mf.demo_pickle()),
          "hex": {"importStep": hexed(mf.demo_pickle(), ops[18][0], ops[19][0]), "callStep": hexed(mf.demo_pickle(), ops[24][0], ops[25][0]),
                  "header": hexed(mf.demo_safetensors(), 8, 8 + n["header_bytes"])},
          "text": {"opening": "Opening the pickle called a function; opening the safetensors file called nothing.",
                   "after19": "After step 19 the pickle has called nothing; the safetensors file has called nothing.",
                   "stopped": "The restricted loader stopped the pickle at step 19 and called nothing. The safetensors file opens as before.",
                   "importLine": "Step 19 of 28 (GLOBAL): IMPORT harmless.ran_on_open. A file could name any function Python can import here.",
                   "callLine": 'Step 25 of 28 (REDUCE): CALL ran_on_open("the demo file was opened"). It records a note and does nothing else.',
                   "stopLine": "Step 19 of 28: the loader meets the import of harmless.ran_on_open, which is not on its list of 5 names. It stops with an error. "
                               "Nothing was called and the numbers did not load.",
                   "notStarted": "Not started. Move the slider or press Next step.",
                   "one": "Functions called: 1 · Notes recorded: 1", "none": "Functions called: 0 · Notes recorded: 0",
                   "safeStep2": 'Step 2 of 3: read the 68-byte header. It is plain text: layer.weight is F32, shape [2, 3], at bytes 0 to 24.',
                   "safeDone": "Step 3 of 3: read those 24 bytes as 6 numbers: 1.5, -0.25, 0.75, 2.0, -1.0, 0.125. Done: nothing was imported or called."}}
    result = run_node(tmp_path, PAGE, fx, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"bad": []}