"""Multi-GPU VRAM budgeting and tensor-split estimation.

llama.cpp's real `--tensor-split` distributes tensor rows across devices; we
approximate it by splitting the **weights** proportionally to each GPU's VRAM
(this is the common case when users want an even load). The KV cache and
compute scratch are placed on a single device (GPU 0 by default, or the GPU
with the most free space) because llama.cpp keeps the cache on the first
device unless manually offloaded.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class GpuBudget:
    vram_bytes: list[int]


@dataclass
class GpuAssignment:
    index: int
    vram_bytes: int
    weight_bytes: float
    kv_compute_bytes: float
    used_bytes: float
    free_bytes: float
    fits: bool


@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
    kv_gpu_index: int = 0


def gpu_split(
    *,
    gpu_vram_bytes: list[int],
    weights_bytes: float,
    kv_compute_bytes: float,
    kv_on_largest: bool = False,
) -> GpuSplitResult:
    """Split weights proportionally to VRAM; place KV+compute on one GPU."""
    if not gpu_vram_bytes:
        return GpuSplitResult()

    total_vram = sum(gpu_vram_bytes)
    if total_vram <= 0:
        return GpuSplitResult(
            assignments=[
                GpuAssignment(i, v, 0.0, 0.0, 0.0, float(v), True)
                for i, v in enumerate(gpu_vram_bytes)
            ],
            total_vram_bytes=total_vram,
        )

    # weight share per gpu proportional to vram
    weight_shares = [weights_bytes * (v / total_vram) for v in gpu_vram_bytes]

    # choose KV host
    if kv_on_largest:
        kv_idx = max(range(len(gpu_vram_bytes)), key=lambda i: gpu_vram_bytes[i])
    else:
        kv_idx = 0

    assignments = []
    for i, vram in enumerate(gpu_vram_bytes):
        w = weight_shares[i]
        kv = kv_compute_bytes if i == kv_idx else 0.0
        used = w + kv
        free = vram - used
        fits = used <= vram
        assignments.append(
            GpuAssignment(
                index=i,
                vram_bytes=vram,
                weight_bytes=w,
                kv_compute_bytes=kv,
                used_bytes=used,
                free_bytes=free,
                fits=fits,
            )
        )

    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_compute_bytes,
        all_fit=all(a.fits for a in assignments),
        kv_gpu_index=kv_idx,
    )


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