"""Keep the merges and the starting word inside ../widget.html equal to the ones here.

    python3 embed_widget.py          rewrite the two data blocks in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

from letters import DEFAULT_LETTER, DEFAULT_WORD, EXAMPLES, MAX_WORD, SEPARATOR, toy_merges

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
CONFIG_BLOCK = re.compile(r'(<script type="application/json" id="ltr-config">)(.*?)(</script>)', re.S)
MERGES_BLOCK = re.compile(r'(<script type="application/json" id="ltr-merges">)(.*?)(</script>)', re.S)


def config():
    return {"word": DEFAULT_WORD, "letter": DEFAULT_LETTER, "maxWord": MAX_WORD, "separator": SEPARATOR,
            "examples": [[word, letter] for word, letter in EXAMPLES.items()]}


def embed(html):
    """html with both data blocks rewritten: the starting state, and the merges the toy tokenizer learns."""
    merges = json.dumps([list(pair) for pair in toy_merges()], separators=(",", ":"))
    html = CONFIG_BLOCK.sub(lambda m: m.group(1) + "\n" + json.dumps(config(), indent=1) + "\n" + m.group(3), html, count=1)
    return MERGES_BLOCK.sub(lambda m: m.group(1) + "\n" + merges + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data blocks now match letters.py and the merges of bpe.py")