"""Keep the numbers and model shape inside ../widget.html equal to quantize.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 quantize as qz

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
DATA_BLOCK = re.compile(r'(<script type="application/json" id="qnt-data">)(.*?)(</script>)', re.S)


def widget_data() -> dict:
    cfg = qz.load_config()
    return {"model": {"name": "Llama 3.1 8B", "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"))},
            "bits": list(qz.BITS), "full_bits": qz.FULL_BITS, "scale_bits": qz.SCALE_BITS, "group": qz.GROUP,
            "sample": {"weights": qz.sample_weights(), "outlier_slot": qz.SAMPLE_OUTLIER_SLOT, "outlier": qz.SAMPLE_OUTLIER,
                       "group": qz.SAMPLE_GROUP},
            "gpus": [[name, int(bw)] for name, bw in qz.BANDWIDTHS.items()]}


def embed(html: str) -> str:
    """html with the data block rewritten from quantize.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 quantize.py")