memory_utils.py

"""Helpers for estimating llama.cpp memory usage for the TUI."""
from __future__ import annotations

import os
import re
import shutil
import subprocess
import textwrap
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

try:
    from huggingface_hub import model_info  # type: ignore
    _HF_OK = True
except Exception:  # pragma: no cover - optional dependency
    model_info = None  # type: ignore
    _HF_OK = False


@dataclass
class GPUInfo:
    index: int
    name: str
    total: Optional[int]
    free: Optional[int]
    uuid: Optional[str] = None
    system_index: Optional[int] = None


@dataclass
class GPUUsage:
    info: GPUInfo
    weights: int
    kv: int


@dataclass
class MemoryProfile:
    source_label: str
    model_label: str
    quant: Optional[str]
    param_count: Optional[float]
    weights_total: Optional[int]
    weights_gpu: int
    weights_cpu: int
    kv_total: int
    ctx_size: Optional[int]
    layers_est: Optional[int]
    gpus: List[GPUUsage]
    cpu_total: Optional[int]
    cpu_available: Optional[int]
    cpu_weights: int
    cpu_kv: int
    warnings: List[str]


_QUANT_PATTERNS: Tuple[Tuple[str, int], ...] = (
    ("Q2_K", 2),
    ("Q3_K_S", 3),
    ("Q3_K_M", 3),
    ("Q3_K_L", 3),
    ("Q3_K", 3),
    ("Q4_K_S", 4),
    ("Q4_K_M", 4),
    ("Q4_K_L", 4),
    ("Q4_K", 4),
    ("Q4_1", 4),
    ("Q4_0", 4),
    ("Q5_K_S", 5),
    ("Q5_K_M", 5),
    ("Q5_K_L", 5),
    ("Q5_K", 5),
    ("Q5_1", 5),
    ("Q5_0", 5),
    ("Q6_K", 6),
    ("Q6_0", 6),
    ("Q8_0", 8),
    ("IQ2", 2),
    ("IQ3", 3),
    ("IQ4", 4),
    ("IQ4_XS", 4),
    ("IQ4_XL", 4),
    ("IQ5", 5),
    ("IQ6", 6),
    ("F16", 16),
    ("FP16", 16),
    ("F32", 32),
    ("FP32", 32),
)


_PARAM_RE = re.compile(r"(\d+(?:\.\d+)?)([bmk])", re.IGNORECASE)
_CPU_SNAPSHOT: Optional[Dict[str, Tuple[int, int]]] = None


def format_bytes(num_bytes: Optional[int]) -> str:
    if num_bytes is None:
        return "?"
    if num_bytes < 0:
        num_bytes = 0
    units = ["B", "KB", "MB", "GB", "TB"]
    value = float(num_bytes)
    for unit in units:
        if value < 1024.0 or unit == units[-1]:
            return f"{value:.2f} {unit}" if unit != "B" else f"{int(value)} {unit}"
        value /= 1024.0
    return f"{value:.2f} PB"


def format_params(params: Optional[float]) -> str:
    if not params:
        return "?"
    if params >= 1e9:
        return f"{params / 1e9:.2f}B"
    if params >= 1e6:
        return f"{params / 1e6:.2f}M"
    if params >= 1e3:
        return f"{params / 1e3:.2f}K"
    return f"{int(params)}"


def _infer_quant_bits(name: str) -> Optional[int]:
    up = name.upper()
    for pattern, bits in _QUANT_PATTERNS:
        if pattern in up:
            return bits
    return None


def _infer_params_from_name(name: str) -> Optional[float]:
    match = _PARAM_RE.search(name)
    if not match:
        return None
    value = float(match.group(1))
    unit = match.group(2).lower()
    if unit == "b":
        return value * 1e9
    if unit == "m":
        return value * 1e6
    if unit == "k":
        return value * 1e3
    return None


def _estimate_layers(param_count: Optional[float]) -> Optional[int]:
    if not param_count or param_count <= 0:
        return None
    # Empirical scaling for decoder-only transformers (LLaMA/Qwen family).
    layers = int(round((param_count / 196608.0) ** (1.0 / 3.0)))
    return max(layers, 1)


def _detect_system_memory() -> Tuple[Optional[int], Optional[int]]:
    meminfo = {}
    try:
        with open("/proc/meminfo", "r", encoding="utf-8") as fh:
            for line in fh:
                if ":" not in line:
                    continue
                key, rest = line.split(":", 1)
                parts = rest.strip().split()
                if not parts:
                    continue
                try:
                    value = int(parts[0]) * 1024
                except ValueError:
                    continue
                meminfo[key] = value
    except FileNotFoundError:  # pragma: no cover - non-Linux fallback
        try:
            import psutil  # type: ignore

            info = psutil.virtual_memory()
            return int(info.total), int(info.available)
        except Exception:
            return None, None
    total = meminfo.get("MemTotal")
    available = meminfo.get("MemAvailable") or meminfo.get("MemFree")
    return total, available


def _read_proc_cpu_snapshot() -> Optional[Dict[str, Tuple[int, int]]]:
    snapshot: Dict[str, Tuple[int, int]] = {}
    try:
        with open("/proc/stat", "r", encoding="utf-8") as fh:
            for line in fh:
                parts = line.split()
                if not parts:
                    continue
                label = parts[0]
                if not label.startswith("cpu"):
                    continue
                if label != "cpu" and not label[3:].isdigit():
                    continue
                try:
                    values = [int(value) for value in parts[1:]]
                except ValueError:
                    continue
                if len(values) < 4:
                    continue
                idle = values[3] + (values[4] if len(values) > 4 else 0)
                total = sum(values)
                snapshot[label] = (total, idle)
    except FileNotFoundError:  # pragma: no cover - non-Linux fallback
        return None
    return snapshot or None


_CPU_SNAPSHOT = _read_proc_cpu_snapshot()


def _detect_cpu_usage() -> Tuple[Optional[float], List[Dict[str, Any]]]:
    global _CPU_SNAPSHOT
    current = _read_proc_cpu_snapshot()
    if current:
        previous = _CPU_SNAPSHOT
        _CPU_SNAPSHOT = current
        if previous:
            def _pct(label: str) -> Optional[float]:
                prev = previous.get(label)
                cur = current.get(label)
                if prev is None or cur is None:
                    return None
                total_delta = cur[0] - prev[0]
                idle_delta = cur[1] - prev[1]
                if total_delta <= 0:
                    return None
                used = max(total_delta - idle_delta, 0)
                return round((used / total_delta) * 100.0, 1)

            overall = _pct("cpu")
            cores: List[Dict[str, Any]] = []
            for label in sorted(
                (name for name in current if name != "cpu"),
                key=lambda name: int(name[3:]),
            ):
                percent = _pct(label)
                cores.append({
                    "index": int(label[3:]),
                    "percent": percent,
                })
            return overall, cores

    try:  # pragma: no cover - non-Linux fallback
        import psutil  # type: ignore

        per_cpu = [float(value) for value in psutil.cpu_percent(interval=None, percpu=True)]
        overall = round(sum(per_cpu) / len(per_cpu), 1) if per_cpu else None
        return overall, [{"index": idx, "percent": round(value, 1)} for idx, value in enumerate(per_cpu)]
    except Exception:
        return None, []


def detect_system_usage() -> Dict[str, Any]:
    cpu_percent, cores = _detect_cpu_usage()
    memory_total, memory_available = _detect_system_memory()
    memory_used = None
    memory_percent = None
    if memory_total is not None and memory_available is not None:
        memory_used = max(memory_total - memory_available, 0)
        if memory_total > 0:
            memory_percent = round((memory_used / memory_total) * 100.0, 1)

    cpu_count_logical = os.cpu_count()
    cpu_count_physical = None
    try:
        import psutil  # type: ignore

        cpu_count_logical = psutil.cpu_count(logical=True) or cpu_count_logical
        cpu_count_physical = psutil.cpu_count(logical=False)
    except Exception:
        cpu_count_physical = None

    try:
        load_1, load_5, load_15 = os.getloadavg()
    except Exception:  # pragma: no cover - non-POSIX fallback
        load_1, load_5, load_15 = None, None, None

    return {
        "cpu_percent": cpu_percent,
        "cpu_count_logical": cpu_count_logical,
        "cpu_count_physical": cpu_count_physical,
        "load_1": load_1,
        "load_5": load_5,
        "load_15": load_15,
        "memory_total": memory_total,
        "memory_available": memory_available,
        "memory_used": memory_used,
        "memory_percent": memory_percent,
        "cores": cores,
    }


def _parse_nvidia_smi_int(value: str) -> Optional[int]:
    cleaned = value.strip()
    if not cleaned or cleaned in {"-", "N/A", "[N/A]", "Not Supported", "[Not Supported]"}:
        return None
    try:
        return int(float(cleaned))
    except ValueError:
        return None


def _normalize_gpu_uuid(value: Any) -> Optional[str]:
    text = str(value or "").strip()
    if not text:
        return None
    if text.startswith("GPU-"):
        text = text[4:]
    return text or None


def _display_gpu_uuid(value: Any) -> Optional[str]:
    normalized = _normalize_gpu_uuid(value)
    if not normalized:
        return None
    return f"GPU-{normalized}"


def _run_nvidia_smi_query(query: str, fields: List[str]) -> List[str]:
    if shutil.which("nvidia-smi") is None:
        return []
    cmd = [
        "nvidia-smi",
        f"--query-{query}={','.join(fields)}",
        "--format=csv,noheader,nounits",
    ]
    try:
        proc = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=2.0,
            check=True,
        )
    except Exception:
        return []
    lines = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
    if len(lines) == 1 and "No running compute processes found" in lines[0]:
        return []
    return lines


def _read_process_display_name(pid: int, fallback: str) -> str:
    try:
        raw = Path(f"/proc/{pid}/cmdline").read_bytes()
        if raw:
            first = raw.split(b"\0", 1)[0].decode("utf-8", errors="replace").strip()
            if first:
                return Path(first).name
    except Exception:
        pass
    cleaned = fallback.strip()
    if not cleaned:
        return f"pid {pid}"
    return Path(cleaned).name or cleaned


def detect_gpu_runtime(managed_processes: Optional[Dict[int, Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
    managed_processes = managed_processes or {}
    gpu_stats_lines = _run_nvidia_smi_query(
        "gpu",
        [
            "index",
            "uuid",
            "name",
            "utilization.gpu",
            "utilization.memory",
            "memory.used",
            "memory.free",
            "memory.total",
        ],
    )
    stats_by_key: Dict[str, Dict[str, Any]] = {}
    for line in gpu_stats_lines:
        parts = [part.strip() for part in line.split(",")]
        if len(parts) < 8:
            continue
        system_index = _parse_nvidia_smi_int(parts[0])
        if system_index is None:
            continue
        uuid_key = _normalize_gpu_uuid(parts[1]) or f"system-index-{system_index}"
        stats_by_key[uuid_key] = {
            "index": system_index,
            "system_index": system_index,
            "uuid": _display_gpu_uuid(parts[1]),
            "name": parts[2],
            "utilization_gpu": _parse_nvidia_smi_int(parts[3]),
            "utilization_memory": _parse_nvidia_smi_int(parts[4]),
            "used": None if _parse_nvidia_smi_int(parts[5]) is None else _parse_nvidia_smi_int(parts[5]) * 1024 * 1024,
            "free": None if _parse_nvidia_smi_int(parts[6]) is None else _parse_nvidia_smi_int(parts[6]) * 1024 * 1024,
            "total": None if _parse_nvidia_smi_int(parts[7]) is None else _parse_nvidia_smi_int(parts[7]) * 1024 * 1024,
            "processes": [],
        }

    gpu_infos = detect_gpus()
    runtime_by_key: Dict[str, Dict[str, Any]] = {}
    for gpu in gpu_infos:
        key = _normalize_gpu_uuid(gpu.uuid) or f"cuda-index-{gpu.index}"
        runtime_by_key[key] = {
            "index": gpu.index,
            "system_index": gpu.system_index,
            "uuid": _display_gpu_uuid(gpu.uuid),
            "name": gpu.name,
            "total": gpu.total,
            "free": gpu.free,
            "used": None if gpu.total is None or gpu.free is None else max(gpu.total - gpu.free, 0),
            "utilization_gpu": None,
            "utilization_memory": None,
            "processes": [],
        }

    for key, stats in stats_by_key.items():
        current = runtime_by_key.get(key)
        if current is None:
            runtime_by_key[key] = dict(stats)
            continue
        current.update({
            "uuid": stats.get("uuid"),
            "system_index": stats.get("system_index"),
            "name": stats.get("name") or current.get("name"),
            "total": stats.get("total") if stats.get("total") is not None else current.get("total"),
            "free": stats.get("free") if stats.get("free") is not None else current.get("free"),
            "used": stats.get("used") if stats.get("used") is not None else current.get("used"),
            "utilization_gpu": stats.get("utilization_gpu"),
            "utilization_memory": stats.get("utilization_memory"),
        })

    process_lines = _run_nvidia_smi_query(
        "compute-apps",
        ["gpu_uuid", "pid", "process_name", "used_gpu_memory"],
    )
    for line in process_lines:
        parts = [part.strip() for part in line.split(",", 3)]
        if len(parts) < 4:
            continue
        key = _normalize_gpu_uuid(parts[0])
        if key is None:
            continue
        gpu_entry = runtime_by_key.get(key)
        if gpu_entry is None:
            continue
        pid = _parse_nvidia_smi_int(parts[1])
        if pid is None:
            continue
        raw_name = parts[2]
        used_memory_mb = _parse_nvidia_smi_int(parts[3])
        used_memory = None if used_memory_mb is None else used_memory_mb * 1024 * 1024
        gpu_total = gpu_entry.get("total")
        extra = managed_processes.get(pid, {})
        process_name = _read_process_display_name(pid, raw_name)
        gpu_entry["processes"].append({
            "pid": pid,
            "process_name": process_name,
            "raw_process_name": raw_name,
            "label": extra.get("label"),
            "kind": extra.get("kind"),
            "status": extra.get("status"),
            "detail": extra.get("detail"),
            "used_memory": used_memory,
            "memory_percent": round((used_memory / gpu_total) * 100.0, 1)
            if used_memory is not None and gpu_total not in (None, 0)
            else None,
        })

    out: List[Dict[str, Any]] = []
    def _sort_key(gpu: Dict[str, Any]) -> tuple[int, int]:
        cuda_index = gpu.get("index")
        system_index = gpu.get("system_index")
        return (
            int(cuda_index) if isinstance(cuda_index, int) else 10_000,
            int(system_index) if isinstance(system_index, int) else 10_000,
        )

    for gpu in sorted(runtime_by_key.values(), key=_sort_key):
        if gpu.get("used") is None and gpu.get("total") is not None and gpu.get("free") is not None:
            gpu["used"] = max(gpu["total"] - gpu["free"], 0)
        gpu["memory_percent"] = round((gpu["used"] / gpu["total"]) * 100.0, 1) if gpu.get("used") is not None and gpu.get("total") not in (None, 0) else None
        gpu["processes"] = sorted(
            gpu.get("processes", []),
            key=lambda proc: (proc.get("used_memory") or 0),
            reverse=True,
        )
        out.append(gpu)
    return out


def detect_gpus() -> List[GPUInfo]:
    try:
        import torch

        if not torch.cuda.is_available():
            return []
        infos: List[GPUInfo] = []
        prev_device: Optional[int] = None
        try:
            if torch.cuda.is_initialized():  # type: ignore[attr-defined]
                prev_device = torch.cuda.current_device()
        except Exception:
            prev_device = None
        for idx in range(torch.cuda.device_count()):
            try:
                props = torch.cuda.get_device_properties(idx)
                name = props.name
                total = getattr(props, "total_memory", None)
                uuid = _normalize_gpu_uuid(getattr(props, "uuid", None))
            except Exception:
                name = f"CUDA {idx}"
                total = None
                uuid = None
            free = None
            try:
                torch.cuda.set_device(idx)
                free, total_runtime = torch.cuda.mem_get_info()
                if total is None:
                    total = int(total_runtime)
            except Exception:
                free = None
            infos.append(
                GPUInfo(
                    index=idx,
                    name=name,
                    total=None if total is None else int(total),
                    free=None if free is None else int(free),
                    uuid=uuid,
                )
            )
        if prev_device is not None:
            try:
                torch.cuda.set_device(prev_device)
            except Exception:
                pass
        return infos
    except Exception:
        return []


def parse_device_list(value: str | None, gpus: Optional[List[GPUInfo]] = None) -> tuple[Optional[List[int]], Optional[str]]:
    """Parse a GPU device list (e.g. '0,1' or 'cuda:0,cuda:1'). Returns (indices, error)."""
    raw = str(value or "").strip()
    if not raw:
        return (None, None)
    lowered = raw.lower()
    if lowered in ("all", "auto", "default"):
        return (None, None)
    if lowered in ("none", "cpu", "off"):
        return ([], None)

    tokens = [tok.strip() for tok in re.split(r"[,\s]+", raw) if tok.strip()]
    if not tokens:
        return (None, "GPU device list cannot be empty.")
    indices: List[int] = []
    for token in tokens:
        cleaned = token.lower()
        for prefix in ("cuda:", "gpu:"):
            if cleaned.startswith(prefix):
                cleaned = cleaned[len(prefix):]
        if not cleaned.isdigit():
            return (None, f"Invalid GPU entry '{token}'. Use a comma-separated list like 0,1.")
        idx = int(cleaned)
        if idx not in indices:
            indices.append(idx)

    if gpus is not None and gpus:
        valid = {gpu.index for gpu in gpus}
        missing = [str(idx) for idx in indices if idx not in valid]
        if missing:
            return (None, f"GPU index(es) not available: {', '.join(missing)}.")
    return (indices, None)


def filter_gpus_by_selection(
    gpus: List[GPUInfo],
    selection: str | None,
) -> tuple[List[GPUInfo], Optional[str], Optional[List[int]]]:
    indices, err = parse_device_list(selection, gpus)
    if indices is None:
        return gpus, err, None
    selected = [gpu for gpu in gpus if gpu.index in indices]
    return selected, err, indices


def _distribute(total: int, count: int, strategy: str) -> List[int]:
    if count <= 0:
        return []
    if total <= 0:
        return [0 for _ in range(count)]
    if strategy == "single" or count == 1:
        return [total] + [0 for _ in range(count - 1)]
    if strategy == "priority":
        main = min(total, int(total * 0.6))
        remaining = total - main
        per_other = remaining // (count - 1)
        parts = [main]
        for i in range(1, count):
            share = per_other
            if i - 1 < remaining - per_other * (count - 1):
                share += 1
            parts.append(share)
        parts[0] += total - sum(parts)
        return parts
    # balanced default
    base = total // count
    remainder = total - base * count
    parts = [base for _ in range(count)]
    for i in range(remainder):
        parts[i % count] += 1
    return parts


def parse_tensor_split(value: str | None) -> tuple[str, Optional[List[float]], Optional[str]]:
    """Parse a llama.cpp --tensor-split string.

    Returns (kind, ratios, error):
      - kind="none": empty/unspecified
      - kind="auto": literal "auto"
      - kind="ratios": ratios is a list of floats (not normalised)
      - kind="invalid": error is a human-friendly message
    """
    raw = str(value or "").strip()
    if not raw:
        return ("none", None, None)
    if raw.lower() == "auto":
        return ("auto", None, None)

    if "," in raw or ";" in raw:
        tokens = [t.strip() for t in raw.replace(";", ",").split(",") if t.strip()]
    elif ":" in raw:
        tokens = [t.strip() for t in raw.split(":") if t.strip()]
    else:
        tokens = [t.strip() for t in raw.split() if t.strip()]

    if not tokens:
        return ("invalid", None, "--tensor-split cannot be empty.")

    ratios: List[float] = []
    for token in tokens:
        if token.endswith("%"):
            token = token[:-1].strip()
        try:
            number = float(token)
        except ValueError:
            return (
                "invalid",
                None,
                f"Invalid --tensor-split '{raw}'. Use e.g. '0.6,0.4', '60,40', or 'auto'.",
            )
        if number < 0:
            return ("invalid", None, f"--tensor-split values must be >= 0 (got {token}).")
        ratios.append(number)

    if sum(ratios) <= 0:
        return ("invalid", None, f"--tensor-split must contain at least one value > 0 (got '{raw}').")

    return ("ratios", ratios, None)


def auto_tensor_split(
    gpus: Optional[List[GPUInfo]] = None,
    *,
    policy: str = "vram",
) -> Optional[str]:
    """Return a comma-separated split string based on the chosen policy, or None if not applicable."""
    if gpus is None:
        gpus = detect_gpus()
    if not gpus or len(gpus) < 2:
        return None
    policy = (policy or "vram").strip().lower()
    if policy == "even":
        totals = [1.0 for _ in gpus]
    else:
        totals = []
        for gpu in gpus:
            if policy == "free":
                value = gpu.free if gpu.free is not None else gpu.total
            else:
                value = gpu.total
            if value is None or value <= 0:
                value = 1.0
            totals.append(float(value))
    denom = sum(totals)
    if denom <= 0:
        base = 100 // len(gpus)
        parts = [base for _ in range(len(gpus))]
        for idx in range(100 - base * len(gpus)):
            parts[idx % len(gpus)] += 1
    else:
        raw = [(total / denom) * 100.0 for total in totals]
        parts = [int(x) for x in raw]
        remainder = 100 - sum(parts)
        order = sorted(((raw[i] - parts[i], i) for i in range(len(gpus))), reverse=True)
        for offset in range(remainder):
            parts[order[offset % len(gpus)][1]] += 1
    total = sum(parts)
    if total != 100 and total > 0:
        parts[-1] += 100 - total
    return ",".join(str(max(p, 0)) for p in parts)


def _distribute_by_ratios(total: int, ratios: List[float], count: int) -> List[int]:
    if count <= 0:
        return []
    if total <= 0:
        return [0 for _ in range(count)]

    weights = list(ratios[:count])
    if len(weights) < count:
        weights.extend([0.0 for _ in range(count - len(weights))])

    denom = sum(weights)
    if denom <= 0:
        return _distribute(total, count, "balanced")

    raw = [total * (w / denom) for w in weights]
    parts = [int(x) for x in raw]
    remainder = total - sum(parts)
    if remainder > 0:
        order = sorted(((raw[i] - parts[i], i) for i in range(count)), reverse=True)
        for offset in range(remainder):
            parts[order[offset % count][1]] += 1

    diff = total - sum(parts)
    if diff != 0 and parts:
        parts[-1] += diff

    return parts


def _resolve_remote_size(model_ref: str, state: dict) -> Optional[int]:
    if not model_ref or ":" not in model_ref or not _HF_OK:
        return None
    cache: Dict[str, Optional[int]] = state.setdefault("_remote_size_cache", {})
    if model_ref in cache:
        return cache[model_ref]
    repo, filename = model_ref.split(":", 1)
    token = (state.get("hf_token") or "").strip() or None
    size: Optional[int] = None
    if model_info is None:
        cache[model_ref] = None
        return None
    try:
        info = model_info(repo, token=token)
        for sibling in getattr(info, "siblings", []):
            if getattr(sibling, "rfilename", None) == filename:
                sz = getattr(sibling, "size", None)
                if sz is not None:
                    size = int(sz)
                    break
    except Exception:
        size = None
    cache[model_ref] = size
    return size


def _model_display_name(source: str, state: dict, model_ref: str) -> Tuple[str, Optional[Path]]:
    models_dir = Path(state.get("models_dir", ".")).expanduser()
    local_choice = (state.get("selected_local_model") or "").strip()
    if source == "local":
        if local_choice:
            candidate = (models_dir / local_choice).expanduser()
            if candidate.exists():
                return local_choice, candidate
        if model_ref:
            path = Path(model_ref).expanduser()
            if path.exists():
                return path.name, path
        return local_choice or (Path(model_ref).name if model_ref else "(no model selected)"), None
    return model_ref or "(no model selected)", None


def _bytes_per_param(bits: Optional[int]) -> Optional[float]:
    if not bits:
        return None
    return bits / 8.0


def estimate_memory_profile(state: dict, *, refresh: bool = False) -> MemoryProfile:
    cache_key = (
        state.get("mode"),
        state.get("model_source"),
        state.get("model_ref"),
        state.get("selected_local_model"),
        state.get("models_dir"),
        state.get("ctx_size"),
        state.get("n_gpu_layers"),
        state.get("tensor_split"),
        state.get("gpu_strategy"),
        state.get("gpu_devices"),
        state.get("auto_split_policy"),
        state.get("hf_token"),
    )
    cache = state.get("_memory_profile_cache")
    if cache and not refresh and cache.get("key") == cache_key:
        profile = cache.get("profile")
        if isinstance(profile, MemoryProfile):
            return profile

    warnings: List[str] = []
    source = state.get("model_source", "local")
    model_ref = str(state.get("model_ref") or "").strip()
    model_label, model_path = _model_display_name(source, state, model_ref)

    quant_bits = _infer_quant_bits(model_label or model_ref)
    quant_tag = None
    if quant_bits is not None:
        for pattern, bits in _QUANT_PATTERNS:
            if bits == quant_bits and pattern in (model_label or model_ref).upper():
                quant_tag = pattern
                break
    param_count = _infer_params_from_name(model_label or model_ref)

    file_size: Optional[int] = None
    if model_path is not None and model_path.exists():
        try:
            file_size = model_path.stat().st_size
        except OSError:
            file_size = None
    if file_size is None and source == "remote":
        remote_size = _resolve_remote_size(model_ref, state)
        if remote_size:
            file_size = remote_size
        elif model_ref:
            warnings.append("Unable to determine remote file size; estimates may be rough.")

    if file_size and (quant_bits is None or param_count is None):
        # Back-compute missing attributes when possible.
        if quant_bits is not None and param_count is None:
            param_count = (file_size * 8.0) / float(quant_bits)
        elif quant_bits is None and param_count is not None and param_count > 0:
            quant_bits = max(int(round((file_size * 8.0) / param_count)), 1)
        elif quant_bits is None and param_count is None:
            param_count = file_size / 2.0  # fallback guess

    bytes_per_param = _bytes_per_param(quant_bits)
    weights_total: Optional[int]
    if param_count and bytes_per_param:
        weights_total = int(param_count * bytes_per_param)
    else:
        weights_total = file_size
    if weights_total is None:
        warnings.append("Model size unknown; memory planner will be limited.")
        weights_total = 0

    layers_est = _estimate_layers(param_count)
    if layers_est is None:
        warnings.append("Could not infer layer count; assuming all layers on GPU when allowed.")
    ctx_value = str(state.get("ctx_size") or "").strip()
    ctx_size: Optional[int]
    if ctx_value:
        try:
            ctx_size = max(int(ctx_value), 1)
        except ValueError:
            warnings.append("Context size is not an integer; using 4096.")
            ctx_size = 4096
    else:
        ctx_size = 4096

    hidden_dim = layers_est * 128 if layers_est else None
    kv_total = 0
    if hidden_dim and ctx_size and layers_est:
        kv_total = int(ctx_size * layers_est * hidden_dim * 4)
    elif ctx_size:
        kv_total = int(ctx_size * 1024 * 16)
        warnings.append("KV cache estimated using generic transformer scaling.")

    n_gpu_layers_raw = str(state.get("n_gpu_layers") or "").strip()
    gpu_ratio = 1.0
    if source == "remote" and not model_ref:
        gpu_ratio = 0.0
    if n_gpu_layers_raw:
        try:
            ngl_value = int(n_gpu_layers_raw)
        except ValueError:
            warnings.append("--n-gpu-layers is not an integer; using automatic placement.")
            ngl_value = 999
    else:
        ngl_value = 999
    if ngl_value <= 0:
        gpu_ratio = 0.0
    elif layers_est and ngl_value not in (999, 1000):
        gpu_ratio = min(max(ngl_value / float(layers_est), 0.0), 1.0)
    elif layers_est is None and ngl_value not in (999, 1000):
        warnings.append("Layer count unknown; treating --n-gpu-layers as full offload.")
        gpu_ratio = 1.0 if ngl_value > 0 else 0.0

    gpus_all = detect_gpus()
    gpus, selection_err, selection_indices = filter_gpus_by_selection(
        gpus_all,
        state.get("gpu_devices"),
    )
    if selection_err:
        warnings.append(selection_err)
    if selection_indices == []:
        warnings.append("No CUDA GPUs selected; weights will remain on system RAM.")
    strategy = state.get("gpu_strategy") or ("cpu" if not gpus else "single")
    if not gpus and strategy != "cpu" and gpu_ratio > 0:
        warnings.append("No CUDA GPUs available; weights will remain on system RAM.")
        gpu_ratio = 0.0
    weights_gpu = int(weights_total * gpu_ratio)
    weights_cpu = max(weights_total - weights_gpu, 0)
    kv_gpu = int(kv_total * gpu_ratio) if kv_total > 0 else 0
    kv_cpu = max(kv_total - kv_gpu, 0)

    if strategy == "cpu" or not gpus:
        weights_gpu = 0
        weights_cpu = weights_total
        kv_gpu = 0
        kv_cpu = kv_total
    weights_gpu = min(weights_gpu, weights_total)
    kv_gpu = min(kv_gpu, kv_total)

    split_raw = str(state.get("tensor_split") or "").strip()
    split_kind, split_ratios, split_err = parse_tensor_split(split_raw)
    if split_kind == "invalid" and split_err:
        warnings.append(split_err)
        split_kind = "none"

    ratio_weights: Optional[List[float]] = None
    if split_kind == "auto" and gpus:
        policy = str(state.get("auto_split_policy") or "vram")
        if policy.strip().lower() == "even":
            ratio_weights = [1.0 for _ in gpus]
        else:
            ratio_weights = []
            for gpu in gpus:
                if policy.strip().lower() == "free":
                    value = gpu.free if gpu.free is not None else gpu.total
                else:
                    value = gpu.total
                if value is None or value <= 0:
                    value = 1.0
                ratio_weights.append(float(value))
    elif split_kind == "ratios" and split_ratios:
        ratio_weights = split_ratios

    if ratio_weights is not None and gpus and strategy != "cpu":
        if len(ratio_weights) != len(gpus):
            warnings.append(
                f"--tensor-split has {len(ratio_weights)} value(s) but {len(gpus)} GPU(s) detected; estimates may be off."
            )
        distribution = _distribute_by_ratios(weights_gpu, ratio_weights, len(gpus))
        kv_distribution = _distribute_by_ratios(kv_gpu, ratio_weights, len(gpus))
    else:
        distribution = _distribute(weights_gpu, len(gpus), strategy)
        kv_distribution = _distribute(kv_gpu, len(gpus), strategy)
    gpu_usages: List[GPUUsage] = []
    for idx, gpu in enumerate(gpus):
        weights_share = distribution[idx] if idx < len(distribution) else 0
        kv_share = kv_distribution[idx] if idx < len(kv_distribution) else 0
        gpu_usages.append(GPUUsage(info=gpu, weights=weights_share, kv=kv_share))

    cpu_total, cpu_available = _detect_system_memory()
    cpu_kv = kv_cpu
    cpu_usage_profile = MemoryProfile(
        source_label="Local" if source == "local" else "Remote",
        model_label=model_label,
        quant=quant_tag,
        param_count=param_count,
        weights_total=weights_total,
        weights_gpu=weights_gpu,
        weights_cpu=weights_cpu,
        kv_total=kv_total,
        ctx_size=ctx_size,
        layers_est=layers_est,
        gpus=gpu_usages,
        cpu_total=cpu_total,
        cpu_available=cpu_available,
        cpu_weights=weights_cpu,
        cpu_kv=cpu_kv,
        warnings=warnings,
    )

    state["_memory_profile_cache"] = {"key": cache_key, "profile": cpu_usage_profile}
    return cpu_usage_profile


def clear_cached_profile(state: dict) -> None:
    state.pop("_memory_profile_cache", None)


def profile_summary_lines(profile: MemoryProfile) -> List[str]:
    lines: List[str] = []
    summary = f"{profile.source_label}: {profile.model_label}"
    bits = []
    if profile.quant:
        bits.append(profile.quant)
    if profile.param_count:
        bits.append(f"≈{format_params(profile.param_count)} params")
    if bits:
        summary += " (" + ", ".join(bits) + ")"
    lines.append(summary)
    if profile.weights_total is not None:
        lines.append(
            f"Weights {format_bytes(profile.weights_total)} → GPU {format_bytes(profile.weights_gpu)} / CPU {format_bytes(profile.weights_cpu)}"
        )
    if profile.kv_total:
        ctx_desc = profile.ctx_size or "auto"
        kv_gpu = sum(gpu.kv for gpu in profile.gpus)
        if kv_gpu or profile.cpu_kv:
            lines.append(
                f"KV cache @ ctx {ctx_desc}: {format_bytes(profile.kv_total)} → GPU {format_bytes(kv_gpu)} / CPU {format_bytes(profile.cpu_kv)}"
            )
        else:
            lines.append(f"KV cache @ ctx {ctx_desc}: {format_bytes(profile.kv_total)}")
    return lines


def warning_lines(profile: MemoryProfile) -> Iterable[str]:
    for warn in profile.warnings:
        yield from textwrap.wrap(warn, 80)