"""Keep the numbers and model shapes inside ../widget.html equal to gpu_split.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 gpu_split as g

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


def widget_data() -> dict:
    models = []
    for key, (name, _) in g.MODELS.items():
        cfg = g.load_config(key)
        models.append({"key": key, "name": name, "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, "counts": list(g.COUNTS), "words": list(g.WORDS), "splits": list(g.SPLITS),
            "links": [{"key": k, "label": label, "bandwidth": bandwidth, "step_us": step_us} for k, (label, bandwidth, step_us) in g.LINKS.items()],
            "gpu": {"gib": g.GPU_GIB, "bandwidth": g.GPU_BANDWIDTH, "flops": g.GPU_FLOPS, "usable": g.USABLE},
            "base_us": g.BASE_US, "all_reduces": g.ALL_REDUCES_PER_LAYER, "word_bytes": g.WORD_BYTES, "chat_tokens": g.CHAT_TOKENS,
            "act_tokens": g.ACTIVATION_TOKENS, "overhead_gib": g.OVERHEAD_GIB, "cache_value_bytes": g.CACHE_VALUE_BYTES}


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