"""Guards for ../widget.html: its data block, its size, that its JavaScript gives the same weights and words as attention.py, and what a reader sees."""
import html
import json
import math
import re
import subprocess

import attention as att
from embed_widget import CONFIG_BLOCK, WIDGET, embed

HTML = WIDGET.read_text(encoding="utf-8")
ROOT = HTML[HTML.index('<div class="attn-root">'):HTML.index("<script")]
TOLERANCE = 1e-12                                   # the JavaScript and Python weights and outputs may differ by at most this much

# Reads the model out of the widget between its marker comments, runs it with Node, and compares the weights and outputs with the numbers
# Python wrote into the fixture (within 1e-12) and the words and percentages as exact strings, for every head and every word.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget, tolerance] = 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="attn-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 {attend: attend, view: view};')(cfg);
const tol = Number(tolerance), bad = [];
let checks = 0, worst = 0;
const near = (a, b) => { const d = Math.abs(a - b); worst = Math.max(worst, d); return d <= tol; };
Object.keys(fx.heads).forEach((name) => {
  const run = api.attend(name), want = fx.heads[name];
  want.weights.forEach((row, i) => row.forEach((w, j) => { checks++; if (!near(run.weights[i][j], w)) bad.push(['weight', name, i, j, run.weights[i][j], w]); }));
  want.outputs.forEach((row, i) => row.forEach((o, c) => { checks++; if (!near(run.outputs[i][c], o)) bad.push(['output', name, i, c, run.outputs[i][c], o]); }));
  want.scores.forEach((row, i) => row.forEach((s, j) => {
    checks++;
    const got = run.scores[i][j];
    if (s === null ? got !== -Infinity : !near(got, s)) bad.push(['score', name, i, j, got, s]);
  }));
  want.views.forEach((v, i) => {
    checks++;
    if (JSON.stringify({ ...api.view(name, i), top: undefined }) !== JSON.stringify({ ...v, top: undefined })) bad.push(['view', name, i, api.view(name, i), v]);
  });
});
console.log(JSON.stringify({ checks, worst, bad: bad.slice(0, 4) }));
process.exit(bad.length ? 1 : 0);
"""

# Runs the page script against a small stand-in for the DOM, presses every word and every head button, 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="attn-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, style: {}, attrs: {}, textContent: '' };
  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['attn-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f) };
  new Function('document', 'window', 'parent', code)(doc, win, page.parent);
  return page;
}
const fire = (e) => (e.listeners.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 pct = (c) => (c === '<1' ? 0.5 : Number(c));
const mix = (td) => { const m = /var\(--green\) (\d+)%/.exec(td.style.background || ''); return m ? Number(m[1]) : NaN; };
function look(page, head, i, what) {                 // compares the page with Python's view of one head and one selected word
  const v = fx.views[head][i], rows = page['attn-bars'].children;
  want(what + ': sentence', page['attn-sentence'].textContent, v.sentence);
  want(what + ': blend', page['attn-blend'].textContent, v.blend);
  want(what + ': scores line', page['attn-scores'].textContent, v.scoresLine);
  want(what + ': word pressed', page['attn-tokens'].children.map((b) => b.attrs['aria-pressed']), fx.tokens.map((t, k) => String(k === i)));
  want(what + ': head pressed', page['attn-heads'].children.map((b) => b.attrs['aria-pressed']), fx.names.map((n) => String(n === head)));
  want(what + ': percent column', rows.map((li) => li.children[2].textContent), v.weights.map((w) => (w === '' ? '' : w + '%')));
  want(what + ': bar widths', rows.map((li) => li.children[1].children[0].style.width), v.widths);
  want(what + ': later words say so', rows.map((li) => [li.attrs['data-later'], li.children[1].children[1].style.display]), v.weights.map((w) => [String(w === ''), w === '' ? '' : 'none']));
  want(what + ': the biggest is marked', rows.map((li) => li.attrs['data-top']), v.weights.map((w, j) => String(j === v.top)));
  const grid = page['attn-grid'].children;
  want(what + ': grid size', grid.length, fx.tokens.length + 1);
  want(what + ': grid numbers', grid.slice(1).map((tr) => tr.children.slice(1).map((td) => td.textContent)), v.grid);
  // shaded only where the word looks; in a row whose weights differ the biggest cell is mixed with more green than the smallest, and no mix passes 50%
  want(what + ': grid shading only where looked, darker for a bigger weight, never above half green', grid.slice(1).map((tr, r) => {
    const cells = tr.children.slice(1), shares = v.grid[r].map((c) => (c === '' ? NaN : pct(c))), seen = shares.map((s, j) => j).filter((j) => !Number.isNaN(shares[j]));
    const mixes = cells.map((td) => ((td.style.background || '') !== '' ? mix(td) : NaN));
    const big = seen.reduce((a, j) => (shares[j] > shares[a] ? j : a), seen[0]), small = seen.reduce((a, j) => (shares[j] < shares[a] ? j : a), seen[0]);
    return [cells.map((td) => (td.style.background || '') !== ''), shares[big] > shares[small] ? mixes[big] > mixes[small] : true, Math.max(...seen.map((j) => mixes[j])) <= 50];
  }), v.grid.map((r) => [r.map((c) => c !== ''), true, true]));
  want(what + ': grid row marked', grid.slice(1).map((tr) => tr.attrs['data-selected']), fx.tokens.map((t, k) => String(k === i)));
}

const first = open();
look(first, fx.opening.head, fx.opening.token, 'opening state');
want('nothing spoken yet', first['attn-live'].textContent, '');
want('word buttons', first['attn-tokens'].children.map((b) => b.textContent), fx.tokens);
want('head buttons', first['attn-heads'].children.map((b) => b.textContent), fx.names.map((n) => fx.labels[n]));
want('grid headers', first['attn-grid'].children[0].children.slice(1).map((th) => th.children[0].textContent), fx.tokens);

// every head with every word, by pressing the buttons
const page = open();
fx.names.forEach((head, h) => {
  fire(page['attn-heads'].children[h]);
  fx.tokens.forEach((t, i) => {
    fire(page['attn-tokens'].children[i]);
    look(page, head, i, head + ' / ' + t);
    want(head + ' / ' + t + ': spoken', page['attn-live'].textContent, fx.views[head][i].sentence + ' ' + fx.views[head][i].blend);
  });
});

// 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' });
send({ type: 'spicyshark:theme', theme: 'light' }, themed.parent);
want('theme back to light', themed.root.attrs, { 'data-theme': 'light' });

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


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 share(text):
    return 0.5 if text == "<1" else float(text)


def test_widget_data_block_matches_the_python_numbers():
    """Guards a toy edit that was not copied into the widget (run embed_widget.py), and an embed that skips a drifted block."""
    assert embed(HTML) == HTML
    drifted = HTML.replace('"opening": {"head": "refers"', '"opening": {"head": "even"', 1).replace('"tokens": ["The", "cat"', '"tokens": ["A", "cat"', 1)
    assert drifted != HTML and embed(drifted) == HTML


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: the parity run covers only the formulas, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    missing = sorted(i for i in lookups if i != "attn-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="attn-config"' in HTML


def test_widget_is_small_and_asks_for_nothing_online():
    assert len(HTML.splitlines()) <= 260 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
    assert "<form" not in HTML and 'type="submit"' not in HTML          # a form cannot submit inside the lesson page's sandboxed frame
    assert "transition" not in HTML and "animation" not in HTML         # nothing moves, so prefers-reduced-motion has nothing to switch off


def test_widget_shows_at_most_five_controls_and_hides_the_rest():
    """Guards the first screen at 390 px: the row of words counts as one control, each head button as one, and the grid sits 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)) + 1 + len(att.HEADS) <= 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=\"attn-grid\"", 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)", ":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", "#f6f8fb", "#05070c", "#00855F", "#009E73"}


def test_widget_formulas_give_the_weights_and_words_of_the_python_toy(tmp_path):
    """Guards the widget's JavaScript drifting from attention.py: the lesson prints Python's numbers and the reader sees the widget's, so
    the two must agree within 1e-12 on every weight, score and output, and exactly on every percentage and sentence, for every head and word."""
    heads = {}
    for name, head in att.HEADS.items():
        run = att.attend(att.VECTORS, head)
        heads[name] = {"weights": run["weights"], "outputs": run["outputs"],
                       "scores": [[None if s == -math.inf else s for s in row] for row in run["scores"]],
                       "views": [att.view(name, i) for i in range(len(att.TOKENS))]}
    result = run_node(tmp_path, PARITY, {"heads": heads}, lambda d: [str(d / "fixture.json"), str(WIDGET), str(TOLERANCE)])
    n, h = len(att.TOKENS), len(att.HEADS)
    assert result["bad"] == [] and result["checks"] == h * (2 * n * n + n * att.D_V + n) and result["worst"] <= TOLERANCE


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: it opens on "it" looking for what it points to with 68% on "cat", every head and every word give
    the bars, sentence, blend line and grid Python computes, later words say they are not looked at, the buttons speak, and the grid figure
    prints the same percentages."""
    opening = att.view("refers", att.TOKENS.index("it"))
    assert opening["sentence"] == "“it” puts 68% of its weight on “cat”." and json.loads(CONFIG_BLOCK.search(HTML).group(2))["opening"] == att.config()["opening"]
    visible = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)                                                                 # what shows without opening anything
    for line in ("A toy: every number is picked by hand, not learned.", "A weight is a share of attention."):
        assert line in visible, line
    assert 'class="attn-sr" id="attn-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT                               # what a screen reader hears
    assert "grid-template-columns: repeat(auto-fit, minmax(var(--min), 1fr))" in HTML and ".attn-root .attn-row button { white-space: nowrap;" in HTML   # an even grid: a button wraps whole, never its label, never stretched
    labels = re.findall(r'class="attn-label"[^>]*>([^<]+)<', ROOT) + list(att.HEAD_LABELS.values())
    assert len(labels) == 2 + len(att.HEADS) and not [t for t in labels if {w.lower() for w in t.split()} & {w.lower() for w in att.TOKENS}]   # no label reads as a word of the sentence ("it")
    lesson, legend = (WIDGET.parent / "lesson.md").read_text(encoding="utf-8"), (WIDGET.parent / "figs" / "two-heads.svg").read_text(encoding="utf-8")
    for key in ("refers", "before"):                                                                                                  # figure 3 and its caption print the widget's own labels
        label = att.HEAD_LABELS[key]
        assert f">{label}</text>" in legend and f"labelled {label}," in legend and f'"{label}"' in lesson, key        # drawn legend, the figure's description, the caption
    named = set(re.findall(r'In the head (?:called )?"([^"]+)"', lesson))                                                             # and so does every sentence that names a head
    assert named and named <= set(att.HEAD_LABELS.values()), named
    grid_svg = (WIDGET.parent / "figs" / "grid.svg").read_text(encoding="utf-8")                                                     # the figure prints the percentages Python computes
    assert [html.unescape(t) for t in re.findall(r'text-anchor="middle"[^>]*>([^<]+)</text>', grid_svg)] == [c for row in att.view("refers", len(att.TOKENS) - 1)["grid"] for c in row if c]
    views = {}
    for name in att.HEADS:
        views[name] = []
        for i in range(len(att.TOKENS)):
            v = att.view(name, i)
            top = -1 if (i > 0 and len(set(v["weights"][:i + 1])) == 1) else max(range(i + 1), key=lambda j: (share(v["weights"][j]), -j))
            v["top"] = top
            v["widths"] = [f"{share(w):g}%" if w != "" else "0%" for w in v["weights"]]
            v["scoresLine"] = (f"Scores before softmax for “{att.TOKENS[i]}”: "
                               + ", ".join(f"{att.TOKENS[j]} {s}" for j, s in enumerate(v["scores"]) if s != "") + ".")
            views[name].append(v)
    ids = sorted(set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):])))
    fixture = {"ids": ids, "tokens": att.TOKENS, "names": list(att.HEADS), "labels": att.HEAD_LABELS, "views": views, "opening": att.config()["opening"]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"heads": len(att.HEADS), "bad": []}