"""Guards for ../widget.html: its data blocks, its size, and that its model gives the numbers toy_lm.py gives."""
import json
import re
import subprocess

import embed_widget
from embed_widget import WIDGET, embed
from prompts import PROMPTS
from toy_lm import (answer, lines_of, load_corpus, model_with_extra, next_word_probs, ranked, smallest_k, tokenize, train,
                    with_temperature)

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

# Reads the model out of the widget between its marker comments, runs it with Node, and compares with the
# numbers Python wrote into the fixture: the same words, the same order, probabilities within 1e-12.
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('/* ---------- the toy model'), html.indexOf('/* ---------- state'));
const api = new Function("'use strict';" + js + '; return {tokenize, linesOf, train, nextWordProbs, withTemperature, ranked, answer, smallestK};')();
const bad = [];
let checks = 0;
fx.tokens.forEach((t) => { checks++; if (JSON.stringify(api.tokenize(t[0])) !== JSON.stringify(t[1])) bad.push(['tokenize', t[0]]); });
fx.cases.forEach((c, i) => {
  const lines = api.linesOf(c.corpus);
  const model = api.train(lines);
  checks++;
  if (JSON.stringify(model.vocab) !== JSON.stringify(c.vocab)) { bad.push(['vocab', i]); return; }
  const probs = api.withTemperature(api.nextWordProbs(api.train(lines.concat(Array(c.k).fill(c.sentence))), c.prompt), c.temp);
  checks++;
  if (probs.size !== c.probs.length) { bad.push(['size', i]); return; }
  c.probs.forEach(([w, p]) => { if (!(Math.abs(probs.get(w) - p) <= 1e-12)) bad.push(['prob', i, w, probs.get(w), p]); });
  checks++;
  if (JSON.stringify(api.ranked(probs).map((e) => e[0])) !== JSON.stringify(c.ranked)) bad.push(['order', i]);
  checks++;
  if (api.answer(probs, c.threshold) !== c.answer) bad.push(['answer', i, api.answer(probs, c.threshold), c.answer]);
  if (c.truth !== null) {
    checks++;
    const k = api.smallestK(lines, c.prompt, c.sentence, c.truth, c.threshold, c.temp, 12);
    if (k !== c.smallest) bad.push(['smallest k', i, k, c.smallest]);
  }
});
console.log(JSON.stringify({ checks, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""

# Corpora a reader might paste: words that are also JavaScript object keys, non-ASCII text, an empty corpus.
CORPORA = {
    "lesson": load_corpus(),
    "odd words": "constructor is __proto__ .\nhasOwnProperty is toString !\nThe café is open, isn't it?\nNo. 2026 is 2026.\n\n\n",
    "one line": "The capital of Peru is Lima.\n",
    "empty": "",
    "one word": "x\n",                     # its only word has probability exactly 1, the edge of the answer-or-abstain check
}
PROMPT_CASES = [(p["prompt"], p["truth"]) for p in PROMPTS] + [
    ("The capital of France is", "Paris"), ("", "The"), ("constructor", "is"), ("the capital of peru is", "Lima"),
    ("is", "a"), ("zzz qqq", "Lima"), ("Café is", "open"),
]
TOKEN_CASES = ["Peru's capital, Lima! Café 2026?", "", "  \t ", "a_b-c", "éè", "end.", "x" * 40, "A.B,C!D?E"]


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


def test_embed_rewrites_a_widget_whose_data_blocks_drifted():
    """Guards embed() skipping a block: a hand-edited corpus line or question must be put back."""
    drifted = HTML.replace("Tokyo.", "Kyoto.", 1).replace('"Lima"', '"Lime"', 1)
    assert drifted != HTML and embed(drifted) == HTML


def test_every_element_the_script_looks_up_exists():
    """Guards a widget that never starts: the parity harness runs only the model, so a misspelt id passed here
    and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    assert lookups and all(f'id="{i}"' in HTML for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in HTML)


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():
    assert len(HTML.encode("utf-8")) < 120 * 1024
    page = embed_widget.CORPUS_BLOCK.sub("", embed_widget.PROMPTS_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_model_gives_the_numbers_of_the_python_code(tmp_path):
    """Guards the widget's JavaScript drifting from toy_lm.py: the lesson prints Python's numbers and the
    reader sees the widget's, so the two must agree for corpora with odd words, unseen prompts and every setting."""
    cases = []
    for name, corpus in CORPORA.items():
        lines = lines_of(corpus)
        vocab = train(lines)["vocab"]
        for prompt, truth in PROMPT_CASES:
            sentence = f"{prompt} {truth}."
            for k in (0, 1, 3):
                for temp in (0.3, 1.0, 1.7):
                    probs = with_temperature(next_word_probs(model_with_extra(lines, sentence, k), prompt), temp)
                    for threshold in (0.0, 0.5, 1.0):
                        cases.append({"name": name, "corpus": corpus, "prompt": prompt, "sentence": sentence, "truth": truth,
                                      "k": k, "temp": temp, "threshold": threshold, "vocab": vocab,
                                      "probs": [[w, p] for w, p in probs.items()], "ranked": [w for w, _ in ranked(probs)],
                                      "answer": answer(probs, threshold),
                                      "smallest": smallest_k(lines, prompt, sentence, truth, threshold, temp, limit=12)})
    fixture = {"tokens": [[t, tokenize(t)] for t in TOKEN_CASES], "cases": cases}
    (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
    result = json.loads(run.stdout)
    assert result["bad"] == [] and result["checks"] > 1000
    assert any(c["smallest"] == 1 for c in cases) and any(c["answer"] == "I don't know" for c in cases)