"""Memory the attention cache needs per token, from a real model configuration.

Python standard library only. The saved file qwen3_8b_config.json is a copy of
https://huggingface.co/Qwen/Qwen3-8B/resolve/main/config.json (Apache-2.0
licence, see the model card). `fetch_config()` downloads the live file.
"""
import json
from pathlib import Path
from urllib.request import urlopen

CONFIG_URL = "https://huggingface.co/Qwen/Qwen3-8B/resolve/main/config.json"
SNAPSHOT = Path(__file__).with_name("qwen3_8b_config.json")
BYTES_PER_VALUE = {"float32": 4, "float16": 2, "bfloat16": 2}


def load_config(path=SNAPSHOT) -> dict:
    return json.loads(Path(path).read_text())


def fetch_config(url=CONFIG_URL) -> dict:
    with urlopen(url, timeout=30) as response:
        return json.load(response)


def kv_bytes_per_token(cfg: dict) -> int:
    """2 (one key, one value) x layers x key/value heads x head size x bytes per value.

    Assumes every layer keeps keys and values for all tokens (full attention).
    Models with sliding-window layers or latent attention store less or differ."""
    head_dim = cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    kv_heads = cfg.get("num_key_value_heads", cfg["num_attention_heads"])
    dtype = cfg.get("torch_dtype") or cfg["dtype"]
    return 2 * cfg["num_hidden_layers"] * kv_heads * head_dim * BYTES_PER_VALUE[dtype]


def kv_bytes(cfg: dict, tokens: int) -> int:
    return kv_bytes_per_token(cfg) * tokens


def kv_bytes_at_max_context(cfg: dict) -> int:
    return kv_bytes(cfg, cfg["max_position_embeddings"])