"""Multimodal projector (mmproj) memory.

A multimodal projector (--mmproj file.gguf) is a separate small GGUF that
holds the vision/audio encoder + adapter weights. It is GPU-offloaded by
default (--no-mmproj-offload disables that). It has no KV cache of its own;
its cost is just its weight tensor bytes plus a small activation scratch we
fold into the compute budget.

We size it by summing the tensor element counts from the mmproj GGUF header
multiplied by each tensor's dtype bits-per-element. The mmproj files in the
wild (e.g. mmproj-F16.gguf / mmproj-BF16.gguf) are almost entirely one dtype,
so a per-dtype sum is accurate.
"""

from __future__ import annotations

from dataclasses import dataclass

# bits per element for the dtypes mmproj files actually use. Sourced from the
# GGML type enum; we only need the projector-relevant subset.
_DTYPE_BPE: dict[int, float] = {
    0: 32.0,   # F32
    1: 16.0,   # F16
    30: 16.0,  # BF16
    2: 4.5,    # Q4_0
    4: 4.0,    # Q4_1 (raw 4)
    8: 8.5,    # Q8_0
}


@dataclass
class Mmproj:
    filename: str = ""
    params: int = 0           # total tensor element count (sum of dims)
    bytes_: float = 0.0        # weight bytes
    offload: bool = True       # --mmproj-offload (default on)
    enabled: bool = False


@dataclass
class MmprojBreakdown:
    bytes_: float = 0.0
    enabled: bool = False
    offload: bool = True


def mmproj_bytes_from_tensors(
    tensor_elems_by_dtype: dict[int, int], params: int | None = None,
) -> float:
    """Sum tensor element counts * bits-per-element / 8 -> weight bytes.

    `tensor_elems_by_dtype` maps GGML dtype id -> total element count for that
    dtype (produced by a header walk). `params` is unused but kept for API
    symmetry with the GGUFMetadata path.
    """
    total_bits = 0.0
    for dtype, count in tensor_elems_by_dtype.items():
        bpe = _DTYPE_BPE.get(dtype)
        if bpe is None:
            # unknown dtype: assume f16 (most mmproj tensors) as a fallback
            bpe = 16.0
        total_bits += count * bpe
    return total_bits / 8.0


def mmproj_bytes_from_meta(meta) -> float:
    """Estimate projector weight bytes from a parsed GGUFMetadata/header.

    Accepts a `GGUFMetadata` (has .raw dict and .params) or a raw metadata dict.
    Uses general.parameter_count when present; otherwise sums tensor elems by
    dtype if a tensor dtype histogram is available in meta['.tensor_dtypes'].
    """
    raw = getattr(meta, "raw", None)
    if raw is None and isinstance(meta, dict):
        raw = meta
    if raw is None:
        return 0.0

    # Prefer an explicit tensor-dtype histogram if the caller attached one.
    hist = raw.get(".tensor_dtype_hist") if isinstance(raw, dict) else None
    if isinstance(hist, dict) and hist:
        return mmproj_bytes_from_tensors(
            {int(k): int(v) for k, v in hist.items()}
        )
    # Fall back to general.parameter_count * a default f16 bpw (2 bytes/elem).
    pc = raw.get("general.parameter_count") if isinstance(raw, dict) else None
    if pc:
        return float(pc) * 2.0
    return 0.0