"""Guards for ../widget.html: its data blocks, its size, and that its tokenizer counts like bpe.py."""
import json
import random
import re
import subprocess
from pathlib import Path

from bpe import NUM_MERGES, encode, load_corpus, train_counted
from cost import decode_positions_read, request_cost, sequential_passes
import embed_widget
from embed_widget import WIDGET, embed
from samples import SAMPLES

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

# Reads the tokenizer and bill code out of the widget between its marker comments, runs it with
# stub elements, and compares with the numbers 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 corpus = html.match(/<script type="text\/plain" id="wtk-corpus">([\s\S]*?)<\/script>/)[1];
const presets = html.match(/<script type="application\/json" id="wtk-presets">([\s\S]*?)<\/script>/)[1];
const js = html.slice(html.indexOf('/* ---------- the tokenizer'), html.indexOf('/* ---------- state'));
const make = new Function('$', 'document', "'use strict';" + js + '; return {train, encode, decode, learned, bill, passes, positions};');
const data = (id) => ({ textContent: id === 'wtk-corpus' ? corpus : presets });
const api = make(data, { getElementById: data });
api.train();
const bad = [];
if (JSON.stringify(api.learned) !== JSON.stringify(fx.merges)) bad.push('merges differ');
let checks = 0;
fx.texts.forEach((t, i) => {
  for (const k of fx.ks) {
    checks++;
    if (JSON.stringify(api.encode(t, k)) !== JSON.stringify(fx.expected[i][String(k)])) bad.push(['encode', t, k]);
  }
  if (api.decode(api.encode(t, fx.ks[fx.ks.length - 1])) !== t) bad.push(['round trip', t]);
});
fx.bills.forEach((b) => {
  const r = api.bill({ nin: b.nin, ncache: b.ncache, nout: b.nout, pin: b.p[0], pcache: b.p[1], pout: b.p[2] });
  if (Math.abs(r.total - Number(b.cost)) > 1e-12) bad.push(['bill', b]);
  if (api.passes(b.nout) !== b.passes) bad.push(['passes', b]);
  if (api.positions(b.nin + b.ncache, b.nout) !== b.positions) bad.push(['positions', b]);
});
console.log(JSON.stringify({ checks, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def test_widget_data_blocks_match_corpus_and_samples():
    """Guards a corpus or sample edit that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_every_element_the_script_looks_up_exists_where_it_looks():
    """Guards a widget that never starts: the parity harness stubs $, so a lookup of an element
    outside .wtk-root (the data blocks sit after it) passed here and threw in the browser."""
    root = HTML[HTML.index('<div class="wtk-root">'):HTML.index("<script")]
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    assert lookups and all(f'id="{i}"' in root for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in root)
    assert "$('wtk-corpus')" not in script and "$('wtk-presets')" not in script


def test_corpus_is_safe_to_embed_in_a_script_element():
    corpus = load_corpus()
    assert "\r" not in corpus and "</script" not in corpus.lower() and "<!--" not in corpus


def test_widget_is_small_and_asks_for_nothing_online():
    """The corpus holds a code sample that mentions fetch(), so look at the page without its data blocks."""
    assert len(HTML.encode("utf-8")) < 120 * 1024
    page = embed_widget.CORPUS_BLOCK.sub("", embed_widget.PRESETS_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


def test_widget_tokenizer_and_bill_match_the_python_code(tmp_path):
    """Guards the widget's JavaScript drifting from bpe.py and cost.py: its counts must be the lesson's."""
    corpus = load_corpus()
    learned = train_counted(corpus)
    merges = [pair for pair, _ in learned]
    assert len(merges) == NUM_MERGES
    rng = random.Random(11)
    alphabet = list("abc the 0123456789,.;\n\t  ") + ["é", "한", "국", "\U0001f600", "\x00", "\x7f", " "]
    texts = list(SAMPLES.values()) + [corpus[i:i + 300] for i in range(0, 3000, 500)]
    texts += ["", " ", "\n\n", "a  b", "end  \n", "한국어 ok"]
    # The page's text box turns CRLF into LF and stops at 4,000 UTF-16 units, but the tokenizer itself must
    # agree with bpe.py on CRLF, combining marks, decomposed Hangul and 4,000 characters as they are.
    texts += ["one\r\ntwo\r\n\r\nthree", "a\r\n", "e\u0301 cafe\u0301", "\u1112\u1161\u11ab\u1100\u1173\u11af", corpus[:4000]]
    texts += ["".join(rng.choice(alphabet) for _ in range(rng.randint(1, 50))) for _ in range(80)]
    ks = [0, 1, 100, 300, NUM_MERGES]
    bills = []
    for nin, ncache, nout in [(500, 1500, 400), (0, 0, 0), (1, 0, 1), (123456, 7, 99999)]:
        p = ["2", "0.2", "10"]
        cost = request_cost(nin, ncache, nout, *p)
        allin = max(1, nin + ncache)
        bills.append({"nin": nin, "ncache": ncache, "nout": nout, "p": [float(x) for x in p], "cost": str(cost),
                      "passes": sequential_passes(allin, nout), "positions": decode_positions_read(allin, nout)})
    fixture = {"merges": [[l, r, n] for (l, r), n in learned], "texts": texts, "ks": ks, "bills": bills,
               "expected": [{str(k): encode(t, merges[:k]) for k in ks} for t in texts]}
    (tmp_path / "fixture.json").write_text(json.dumps(fixture), 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)["checks"] == len(texts) * len(ks) >= 400