"""Multi-GPU VRAM budgeting, modelled on llama.cpp's real --split-mode.

llama.cpp distributes work across GPUs in one of three split modes
(see docs/multi-gpu.md):

  layer  (default, pipeline-parallel): each GPU holds a contiguous slice of
          *layers*. The KV cache for layer l lives on the GPU that owns layer
          l, so KV is distributed across GPUs in the same proportion as the
          weights — it is NOT piled on GPU 0.

  tensor (experimental, tensor-parallel): splits both weights *and* KV across
          the participating GPUs, either evenly or by --tensor-split
          proportions. Requires flash-attn on and a non-quantized KV cache
          (f16/bf16/f32); quantized KV is a hard error upstream.

  none   : everything on --main-gpu (default 0); the other GPUs are unused.

We support mixed-GPU setups of differing size (e.g. a 24 GB Radeon + a 16 GB
unified-memory iGPU). When the caller passes explicit --tensor-split ratios we
use those; otherwise we split proportionally to each GPU's VRAM (layer mode) or
to the given ratios (tensor mode). Unified-memory devices (system RAM acting
as VRAM, e.g. AMD Strix Halo / Apple SoCs) are flagged so callers can warn that
the estimate is fuzzy — that memory is shared with the OS.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class GpuSpec:
    """A single device: its VRAM and whether it's unified (shared system RAM)."""
    vram_gb: float
    name: str = ""
    is_unified: bool = False

    @property
    def vram_bytes(self) -> int:
        return int(self.vram_gb * (1 << 30))


@dataclass
class GpuBudget:
    vram_bytes: list[int]


@dataclass
class GpuAssignment:
    index: int
    vram_bytes: int
    name: str = ""
    is_unified: bool = False
    weight_bytes: float = 0.0
    kv_compute_bytes: float = 0.0
    used_bytes: float = 0.0
    free_bytes: float = 0.0
    fits: bool = True
    role: str = ""  # human note: "weights+KV", "weights", "all (main)", ...


@dataclass
class GpuSplitResult:
    assignments: list[GpuAssignment] = field(default_factory=list)
    total_vram_bytes: int = 0
    total_used_bytes: float = 0.0
    total_weights_bytes: float = 0.0
    total_kv_compute_bytes: float = 0.0
    all_fit: bool = False
    split_mode: str = "layer"
    main_gpu_index: int = 0
    warnings: list[str] = field(default_factory=list)


def _normalize_shares(ratios: list[float], n: int) -> list[float]:
    """Turn a list of split ratios into n normalized shares summing to 1.

    None/empty -> equal shares. Otherwise pad/truncate to n and normalize.
    """
    if not ratios:
        s = [1.0] * n
    else:
        s = [float(r) for r in ratios[:n]]
        while len(s) < n:
            s.append(0.0)
    tot = sum(s)
    if tot <= 0:
        s = [1.0] * n
        tot = float(n)
    return [x / tot for x in s]


def gpu_split(
    *,
    gpu_specs: list[GpuSpec] | None = None,
    gpu_vram_bytes: list[int] | None = None,
    weights_bytes: float,
    kv_bytes: float,
    scratch_bytes: float = 0.0,
    split_mode: str = "layer",
    main_gpu: int = 0,
    tensor_split: list[float] | None = None,
    cache_dtype_quantized: bool = False,
) -> GpuSplitResult:
    """Distribute weights, KV, and scratch across GPUs for a split mode.

    `gpu_specs` is the preferred input; `gpu_vram_bytes` is accepted for
    backward compatibility (each treated as a discrete, non-unified GPU).
    `kv_bytes` is the model KV cache; `scratch_bytes` is the compute/activation
    buffer, which is placed on the main GPU (the one running the active batch).
    """
    # Normalize inputs into a list of (vram_bytes, is_unified, name).
    if gpu_specs is not None:
        gpus = [(g.vram_bytes, g.is_unified, g.name) for g in gpu_specs]
    elif gpu_vram_bytes is not None:
        gpus = [(int(v), False, "") for v in gpu_vram_bytes]
    else:
        return GpuSplitResult()

    if not gpus:
        return GpuSplitResult()

    warns: list[str] = []
    if any(u for _, u, _ in gpus):
        warns.append(
            "One or more devices is unified memory (shared system RAM); the "
            "VRAM budget is fuzzy because that memory is also used by the OS."
        )

    n = len(gpus)
    total_vram = sum(v for v, _, _ in gpus)
    if main_gpu < 0 or main_gpu >= n:
        main_gpu = 0

    mode = split_mode if split_mode in ("layer", "tensor", "none") else "layer"

    # --- split mode: none -------------------------------------------------
    if mode == "none":
        v, uni, nm = gpus[main_gpu]
        used = weights_bytes + kv_bytes + scratch_bytes
        assignments = []
        for i, (vr, ur, nr) in enumerate(gpus):
            if i == main_gpu:
                assignments.append(GpuAssignment(
                    index=i, vram_bytes=vr, name=nr, is_unified=ur,
                    weight_bytes=weights_bytes, kv_compute_bytes=kv_bytes,
                    used_bytes=used, free_bytes=vr - used, fits=used <= vr,
                    role="all (main)",
                ))
            else:
                assignments.append(GpuAssignment(
                    index=i, vram_bytes=vr, name=nr, is_unified=ur,
                    used_bytes=0.0, free_bytes=float(vr), fits=True,
                    role="unused",
                ))
        return GpuSplitResult(
            assignments=assignments, total_vram_bytes=total_vram,
            total_used_bytes=used,
            total_weights_bytes=weights_bytes,
            total_kv_compute_bytes=kv_bytes,
            all_fit=all(a.fits for a in assignments),
            split_mode=mode, main_gpu_index=main_gpu, warnings=warns,
        )

    # --- split mode: tensor (experimental) -------------------------------
    if mode == "tensor":
        if cache_dtype_quantized:
            warns.append(
                "--split-mode tensor requires a non-quantized KV cache "
                "(f16/bf16/f32); quantized KV is a hard error in llama.cpp."
            )
        shares = _normalize_shares(tensor_split or [], n)
        # tensor mode splits weights AND KV across all GPUs by the shares;
        # scratch rides on the main GPU.
        assignments = []
        for i, (vr, ur, nr) in enumerate(gpus):
            w = weights_bytes * shares[i]
            kv = kv_bytes * shares[i]
            sc = scratch_bytes if i == main_gpu else 0.0
            used = w + kv + sc
            assignments.append(GpuAssignment(
                index=i, vram_bytes=vr, name=nr, is_unified=ur,
                weight_bytes=w, kv_compute_bytes=kv,
                used_bytes=used, free_bytes=vr - used, fits=used <= vr,
                role="weights+KV" + ("+compute" if i == main_gpu else ""),
            ))
        return GpuSplitResult(
            assignments=assignments, total_vram_bytes=total_vram,
            total_used_bytes=sum(a.used_bytes for a in assignments),
            total_weights_bytes=weights_bytes,
            total_kv_compute_bytes=kv_bytes,
            all_fit=all(a.fits for a in assignments),
            split_mode=mode, main_gpu_index=main_gpu, warnings=warns,
        )

    # --- split mode: layer (default, pipeline-parallel) ------------------
    # Default split is proportional to each GPU's VRAM (llama.cpp's auto-split).
    # If the user gave explicit --tensor-split ratios, use those as the weight
    # proportions instead.
    if tensor_split:
        shares = _normalize_shares(tensor_split, n)
    else:
        if total_vram <= 0:
            shares = [1.0 / n] * n
        else:
            shares = [v / total_vram for v, _, _ in gpus]
    assignments = []
    for i, (vr, ur, nr) in enumerate(gpus):
        w = weights_bytes * shares[i]
        # KV follows the layer that owns it -> distributed like the weights.
        kv = kv_bytes * shares[i]
        sc = scratch_bytes if i == main_gpu else 0.0
        used = w + kv + sc
        assignments.append(GpuAssignment(
            index=i, vram_bytes=vr, name=nr, is_unified=ur,
            weight_bytes=w, kv_compute_bytes=kv,
            used_bytes=used, free_bytes=vr - used, fits=used <= vr,
            role="weights+KV" + ("+compute" if i == main_gpu else ""),
        ))
    return GpuSplitResult(
        assignments=assignments, total_vram_bytes=total_vram,
        total_used_bytes=sum(a.used_bytes for a in assignments),
        total_weights_bytes=weights_bytes,
        total_kv_compute_bytes=kv_bytes,
        all_fit=all(a.fits for a in assignments),
        split_mode=mode, main_gpu_index=main_gpu, warnings=warns,
    )


def fit_gpus(result: GpuSplitResult) -> bool:
    return result.all_fit