"""Keep the numbers and model shapes inside ../widget.html equal to memory_calc.py.

    python3 embed_widget.py          rewrite the data block 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

import memory_calc as mc

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
DATA_BLOCK = re.compile(r'(<script type="application/json" id="wmm-data">)(.*?)(</script>)', re.S)
SHORT = {"1b": "1B", "8b": "8B", "70b": "70B", "405b": "405B"}


def widget_data() -> dict:
    models = []
    for key, (name, _) in mc.MODELS.items():
        cfg = mc.load_config(key)
        models.append({"key": key, "name": name, "short": SHORT[key], "layers": cfg["num_hidden_layers"],
                       "hidden": cfg["hidden_size"], "ffn": cfg["intermediate_size"], "heads": cfg["num_attention_heads"],
                       "kv_heads": cfg["num_key_value_heads"], "head_dim": cfg.get("head_dim"), "vocab": cfg["vocab_size"],
                       "tied": bool(cfg.get("tie_word_embeddings"))})
    return {"models": models, "bits": list(mc.BITS), "contexts": list(mc.CONTEXTS), "chats": list(mc.CHATS),
            "full_bits": mc.FULL_BITS, "gpus": [list(g) for g in mc.GPUS], "usable": mc.USABLE, "overhead_gib": mc.OVERHEAD_GIB,
            "act_tokens": mc.ACTIVATION_TOKENS, "act_bytes": mc.ACTIVATION_BYTES, "kv_value_bytes": mc.KV_VALUE_BYTES}


def embed(html: str) -> str:
    """html with the data block rewritten from memory_calc.py."""
    block = json.dumps(widget_data(), indent=1)
    return DATA_BLOCK.sub(lambda m: m.group(1) + "\n" + block + "\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 block now matches memory_calc.py")