"""What a model needs in memory, derived from one GGUF header.

The weights split into three parts that llama.cpp places differently:

* ``token_embd`` -- read on the CPU side, not offloaded with ``-ngl``;
* ``output`` (the LM head, or the tied embedding standing in for it) -- offloaded
  together with the last layer;
* the repeating blocks ``blk.N.*`` -- offloaded one by one with ``-ngl``, and within
  a MoE block the ``*_exps`` expert tensors, which ``--cpu-moe`` / ``--n-cpu-moe``
  keep in RAM.

The KV cache is per layer and depends on the attention kind: full attention grows
with the context, sliding-window layers stop at the window, latent attention (MLA)
stores one compressed vector, and recurrent layers in hybrids store none.

Geometry does not depend on the quant, so one header describes every quant of a
model; byte counts are scaled to each file's size.
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass, field

from .gguf_header import GGUFHeader

_BLK = re.compile(r"^blk\.(\d+)\.")

# Sliding-window layout for architectures whose header gives a window but no per-layer
# pattern: a global layer every N layers, the rest local. Unknown architectures with a
# window are treated as all-global, which over- rather than under-estimates.
SERVER_SLOTS = 4   # llama-server's default --parallel (auto), with a unified KV cache
UBATCH = 512       # default -ub
SWA_GLOBAL_EVERY = {"gemma2": 2, "gemma3": 6, "gemma3n": 5, "gpt-oss": 2, "cohere2": 4, "exaone4": 4}


@dataclass
class Geometry:
    arch: str | None
    n_layers: int
    context_length: int | None
    params: int | None
    moe: bool
    expert_count: int | None
    expert_used: int | None
    # bytes, for the file the header came from (all shards)
    file_bytes: int
    embd_bytes: int
    output_bytes: int
    block_bytes: int
    expert_bytes: int
    # KV: one entry per layer that has a cache: (bytes per token at f16, window or None)
    kv_layers: list[tuple[float, int | None]] = field(default_factory=list)
    notes: list[str] = field(default_factory=list)
    # weights (not bytes) in the embedding and output matrices, so another quant's split
    # can be worked out from its own bits per weight rather than this file's
    embd_params: int | None = None
    output_params: int | None = None  # None with tied embeddings

    # -- scaled to another quant of the same model
    def scaled(self, file_bytes: int, params: float | None = None) -> "Geometry":
        """This geometry for another quant of the same model.

        Scaling every part by the file-size ratio is wrong for the embedding and output
        matrices: a 1-bit file keeps them near 4 and 6 bits, a Q8_0 file at 8.5. Scaled from
        the smallest quant, the layers (which go to the GPU) came out up to 0.4 GB short.
        When the weight counts are known, those two are sized at the bits quantizers give
        them and the layers get the rest."""
        if not self.file_bytes:
            return self
        k = file_bytes / self.file_bytes
        g = Geometry(**{**self.__dict__})
        g.kv_layers = list(self.kv_layers)
        g.notes = list(self.notes)
        g.file_bytes = file_bytes
        g.embd_bytes = int(self.embd_bytes * k)
        g.output_bytes = int(self.output_bytes * k)
        g.block_bytes = int(self.block_bytes * k)
        g.expert_bytes = int(self.expert_bytes * k)
        p = self.params or params
        if p and self.embd_params:
            avg = file_bytes * 8 / p
            tied = self.output_params is None
            out_bpw = avg if avg > 9 else min(max(avg, 6.5625), 8.5)  # Q6_K for most K-quants, Q8_0 at 8-bit
            embd = self.embd_params * (out_bpw if tied else avg) / 8  # tied: quantized like an output layer
            output = embd if tied else self.output_params * out_bpw / 8
            layers = file_bytes - embd - (0 if tied else output)
            if layers > 0.5 * file_bytes:
                share = self.expert_bytes / self.block_bytes if self.block_bytes else 0.0
                g.embd_bytes, g.output_bytes = int(embd), int(output)
                g.block_bytes, g.expert_bytes = int(layers), int(layers * share)
        return g

    @property
    def layer_bytes(self) -> float:
        return self.block_bytes / max(self.n_layers, 1)

    @property
    def layer_expert_bytes(self) -> float:
        return self.expert_bytes / max(self.n_layers, 1)

    def kv_bytes(self, context: int, kv_bytes_per_elem: float = 2.0) -> float:
        """KV cache for ``context`` tokens; f16 is 2 bytes per element.

        Sliding-window layers hold more than one window: llama-server runs 4 slots on a
        unified cache by default, and llama.cpp sizes those layers for every slot's window
        plus a batch (n_swa * n_seq_max + n_ubatch, padded to 256)."""
        total = 0.0
        for per_token_f16, window in self.kv_layers:
            tokens = min(context, -(-(window * SERVER_SLOTS + UBATCH) // 256) * 256) if window else context
            total += per_token_f16 / 2.0 * kv_bytes_per_elem * tokens
        return total

    def kv_bytes_per_layer(self, context: int, kv_bytes_per_elem: float = 2.0) -> float:
        return self.kv_bytes(context, kv_bytes_per_elem) / max(self.n_layers, 1)

    def to_dict(self) -> dict:
        return {
            "arch": self.arch, "n_layers": self.n_layers, "context_length": self.context_length,
            "params": self.params, "moe": self.moe, "expert_count": self.expert_count,
            "expert_used": self.expert_used, "kv_layers": len(self.kv_layers),
            "swa_layers": sum(1 for _, w in self.kv_layers if w), "notes": self.notes,
        }


def _per_layer(value, n: int, default=None) -> list:
    if isinstance(value, list):
        vals = list(value)[:n]
        return vals + [vals[-1] if vals else default] * (n - len(vals))
    return [value if value is not None else default] * n


def from_header(h: GGUFHeader, file_bytes: int | None = None) -> Geometry:
    md = h.metadata
    arch = h.arch or ""
    get = h.arch_key
    n = int(get("block_count") or 0)
    nextn = int(get("nextn_predict_layers") or 0)

    # -- weight split, from the tensor table of the shard we read
    embd = sum(b for name, b in h.tensors if name.startswith("token_embd."))
    output = sum(b for name, b in h.tensors if name.startswith("output.") and not name.startswith("output_norm"))
    embd_params = sum(v for name, v in h.elements.items() if name.startswith("token_embd.")) or None
    output_params = sum(v for name, v in h.elements.items()
                        if name.startswith("output.") and not name.startswith("output_norm")) or None
    table_params = sum(h.elements.values())
    blocks = sum(b for name, b in h.tensors if _BLK.match(name))
    experts = sum(b for name, b in h.tensors if _BLK.match(name) and "_exps." in name)
    other = sum(b for _, b in h.tensors) - embd - output - blocks
    notes = []
    if not output and embd:
        output = embd  # tied embeddings: llama.cpp puts a copy on the GPU as the output layer
        notes.append("tied embeddings")

    table_bytes = sum(b for _, b in h.tensors)
    split_count = int(md.get("split.count") or 1)
    total = file_bytes or table_bytes
    if not blocks and total:
        # The shard read lists no layer tensors: some splits keep only metadata in the
        # first file. Split the bytes by parameter counts derived from the metadata.
        embd, output, blocks, experts, other = _split_by_params(h, total)
        notes.append("layer sizes estimated from the model's shape")
        embd_params = output_params = table_params = None
    elif split_count > 1 and table_bytes and total > table_bytes:
        # Only the first shard's tensors are listed. Embeddings and output live in it
        # (both are written first); scale the blocks to the rest of the file, keeping
        # the expert share this shard shows.
        expert_share = experts / blocks if blocks else 0.0
        blocks = max(total - embd - output - other, 0)
        experts = int(blocks * expert_share)
        notes.append("split file: layer sizes extrapolated from its first shard")
        table_params = None  # the other shards' weights aren't listed
    elif total and table_bytes and abs(total - table_bytes) > 0.02 * total:
        k = total / table_bytes
        embd, output, blocks, experts = (int(x * k) for x in (embd, output, blocks, experts))

    # -- KV cache geometry
    head_count = _per_layer(get("attention.head_count"), n, 1)
    kv_heads = _per_layer(get("attention.head_count_kv"), n, None)
    kv_heads = [kv if kv is not None else hc for kv, hc in zip(kv_heads, head_count)]
    embed = int(get("embedding_length") or 0)
    key_len = get("attention.key_length")
    val_len = get("attention.value_length")
    key_len_swa = get("attention.key_length_swa", key_len)
    val_len_swa = get("attention.value_length_swa", val_len)
    window = get("attention.sliding_window")
    pattern = get("attention.sliding_window_pattern")
    interval = get("full_attention_interval")

    kv_layers: list[tuple[float, int | None]] = []
    kv_lora = get("attention.kv_lora_rank")
    if kv_lora:  # MLA: one compressed latent + the RoPE key per token per layer
        rope = int(get("rope.dimension_count") or 64)
        per = (int(kv_lora) + rope) * 2.0
        kv_layers = [(per, None)] * max(n - nextn, 0)
        notes.append("latent attention (MLA)")
    else:
        for i in range(max(n - nextn, 0)):
            heads = int(kv_heads[i] or 0)
            if heads == 0:
                continue  # recurrent / convolution layer in a hybrid: no KV cache
            if interval and (i + 1) % int(interval) != 0:
                continue  # hybrid with full attention every `interval` layers
            is_swa = False
            if window:
                if isinstance(pattern, list) and i < len(pattern):
                    is_swa = bool(pattern[i])
                elif arch in SWA_GLOBAL_EVERY:
                    is_swa = (i + 1) % SWA_GLOBAL_EVERY[arch] != 0
            hc = int(head_count[i] or 1)
            k_len = int((key_len_swa if is_swa else key_len) or (embed // hc if hc else 0))
            v_len = int((val_len_swa if is_swa else val_len) or k_len)
            kv_layers.append((heads * (k_len + v_len) * 2.0, int(window) if is_swa else None))
        if interval:
            notes.append(f"hybrid: attention every {int(interval)} layers")
        if any(w for _, w in kv_layers):
            notes.append(f"sliding window {int(window)} on {sum(1 for _, w in kv_layers if w)} layers")

    expert_count = get("expert_count")
    return Geometry(
        arch=arch or None,
        n_layers=n,
        context_length=int(get("context_length")) if get("context_length") else None,
        params=table_params,
        moe=bool(expert_count and int(expert_count) > 1),
        expert_count=int(expert_count) if expert_count else None,
        expert_used=int(get("expert_used_count")) if get("expert_used_count") else None,
        file_bytes=int(total),
        embd_bytes=int(embd),
        output_bytes=int(output),
        block_bytes=int(blocks + max(other, 0)),
        expert_bytes=int(experts),
        kv_layers=kv_layers,
        notes=notes,
        embd_params=embd_params,
        output_params=output_params if (output_params and "tied embeddings" not in notes) else None,
    )


def _split_by_params(h: GGUFHeader, total: int) -> tuple[int, int, int, int, int]:
    """(embd, output, blocks, experts, other) bytes from the header's shape alone."""
    get = h.arch_key
    n = int(get("block_count") or 1)
    d = int(get("embedding_length") or 4096)
    vocab_meta = h.metadata.get("tokenizer.ggml.tokens")
    vocab = int(get("vocab_size") or (vocab_meta.get("array_len") if isinstance(vocab_meta, dict) else 0) or 150_000)
    heads = get("attention.head_count")
    heads = int(max(heads) if isinstance(heads, list) else heads or 32)
    ff = get("feed_forward_length")
    ff = int(max(ff) if isinstance(ff, list) else ff or 4 * d)
    n_exp = int(get("expert_count") or 0)
    exp_ff = int(get("expert_feed_forward_length") or ff)
    shared = int(get("expert_shared_count") or 0)
    dense_lead = int(get("leading_dense_block_count") or 0)
    attn = 4 * d * d
    if n_exp > 1:
        moe_layers = n - dense_lead
        expert_p = moe_layers * n_exp * 3 * d * exp_ff
        rest_p = n * attn + moe_layers * shared * 3 * d * exp_ff + dense_lead * 3 * d * ff
    else:
        expert_p = 0
        rest_p = n * (attn + 3 * d * ff)
    embd_p = vocab * d
    all_p = expert_p + rest_p + 2 * embd_p
    per_param = total / all_p if all_p else 0
    embd = int(embd_p * per_param)
    return embd, embd, int((expert_p + rest_p) * per_param), int(expert_p * per_param), 0


def estimate_without_header(params: float, bpw: float, file_bytes: int | None, context_length: int | None,
                            moe: bool = False) -> Geometry:
    """A rough geometry from the parameter count alone, for models whose header has not
    been read yet. Layer count and KV size follow typical dense-transformer ratios; the
    result is flagged so the UI can say it is an estimate."""
    b = params / 1e9
    n = max(16, int(round(8 * math.log2(max(b, 1)) + 20)))
    total = file_bytes or int(params * bpw / 8)
    embd = int(total * (0.08 if b < 4 else 0.03))
    kv_per_token = 2.0 * 2 * 8 * 128  # 8 KV heads x head dim 128, K+V, f16, per layer
    return Geometry(arch=None, n_layers=n, context_length=context_length, params=int(params), moe=moe,
                    expert_count=None, expert_used=None, file_bytes=total, embd_bytes=embd, output_bytes=embd,
                    block_bytes=max(total - 2 * embd, 0), expert_bytes=int(total * 0.9) if moe else 0,
                    kv_layers=[(kv_per_token, None)] * n, notes=["estimated from parameter count"])