#!/usr/bin/env python3
"""Estimate llama.cpp memory for a model, given a model id.

A "model id" is either:
  - a local .gguf file path, or
  - a Hugging Face repo id (e.g. "bartowski/Qwen2.5-7B-Instruct-GGUF"),
    optionally with a llama.cpp router-style quant tag after ':'
    (e.g. "org/model:Q6_K_P") to select one variant.

It reads only the GGUF header (metadata + tensor table), so it works on
multi-GB models without downloading them.

Source preference for an HF repo id:
  1. The local HF Hub cache (the same layout llama.cpp's own downloader
     uses, see common/hf-cache.cpp): $LLAMA_CACHE / $HF_HUB_CACHE /
     $HUGGINGFACE_HUB_CACHE / $HF_HOME/hub / $XDG_CACHE_HOME/huggingface/hub
     / ~/.cache/huggingface/hub, under models--<org>--<repo>/snapshots/<commit>/.
     Files already there are read straight from disk (offline, instant).
  2. Otherwise, HTTP Range requests against huggingface.co (header only).
  Pass --online to force the network path and ignore the cache.

Numbers it reports (option 5 from the conversation):
  weights   : bytes on disk for the tensor data section
              (= sum over parts of (part_size - data_offset); exact)
  kv cache  : (n_embd_head_k*n_head_kv + n_embd_head_v*n_head_kv)
              * n_ctx * n_layer * bpp(dtype)/8 * n_seqs
              matches src/llama-kv-cache.cpp (n_embd_k_gqa = head_dim*n_head_kv)
  compute   : transient scratch, noted qualitatively (depends on ubatch)

By default prints a single integer: estimated resident memory (weights + kv)
in MiB. Pass -v/--verbose for the full breakdown.

Stdlib only. Uses HF_TOKEN env var or --token for private repos.
"""

import argparse
import json
import os
import pathlib
import re
import struct
import sys
import urllib.request

GGUF_MAGIC = 0x46554747  # "GGUF"

# GGUFValueType ids (gguf-py/gguf/constants.py)
VT_UINT8, VT_INT8, VT_UINT16, VT_INT16 = 0, 1, 2, 3
VT_UINT32, VT_INT32, VT_FLOAT32, VT_BOOL = 4, 5, 6, 7
VT_STRING, VT_ARRAY, VT_UINT64, VT_INT64, VT_FLOAT64 = 8, 9, 10, 11, 12

# bits per element for cache dtypes (ggml/src/ggml.c type_traits + ggml-common.h)
CACHE_BPP = {
    "f32": 32, "f16": 16, "bf16": 16,
    "q8_0": 68, "q8_1": 72, "q4_0": 36, "q4_1": 40,
    "q5_0": 44, "q5_1": 48, "q4_k": 36, "q5_k": 44,
    "q6_k": 52.5, "q8_k": 73, "iq4_nl": 36, "mxfp4": 36,
}

# archs that are recurrent, not KV-cache transformers
RECURRENT_ARCHS = {"mamba", "mamba2", "rwkv6", "rwkv7", "griffin", "dsv4"}

DEFAULT_ALIGN = 32
MiB = 1024 * 1024
GiB = 1024 * 1024 * 1024

# GGMLQuantizationType (ggml_type) - used for tensor.type
QUANT_NAMES = {
    0: "f32", 1: "f16", 2: "q4_0", 3: "q4_1", 6: "q5_0", 7: "q5_1", 8: "q8_0",
    9: "q8_1", 10: "q2_k", 11: "q3_k", 12: "q4_k", 13: "q5_k", 14: "q6_k",
    15: "q8_k", 16: "iq2_xxs", 17: "iq2_xs", 18: "iq3_xxs", 19: "iq1_s",
    20: "iq4_nl", 21: "iq3_s", 22: "iq2_s", 23: "iq4_xs", 24: "i8", 25: "i16",
    26: "i32", 27: "i64", 28: "f64", 29: "iq1_m", 30: "bf16", 34: "tq1_0",
    35: "tq2_0", 39: "mxfp4", 40: "nvfp4", 41: "q1_0", 42: "q2_0",
}

# LlamaFileType (llama_ftype) - used for general.file_type; NOT the same enum
FTYPE_NAMES = {
    0: "all_f32", 1: "mostly_f16", 2: "mostly_q4_0", 3: "mostly_q4_1",
    7: "mostly_q8_0", 8: "mostly_q5_0", 9: "mostly_q5_1", 10: "mostly_q2_k",
    11: "mostly_q3_k_s", 12: "mostly_q3_k_m", 13: "mostly_q3_k_l",
    14: "mostly_q4_k_s", 15: "mostly_q4_k_m", 16: "mostly_q5_k_s",
    17: "mostly_q5_k_m", 18: "mostly_q6_k", 19: "mostly_iq2_xxs",
    20: "mostly_iq2_xs", 21: "mostly_q2_k_s", 22: "mostly_iq3_xs",
    23: "mostly_iq3_xxs", 24: "mostly_iq1_s", 25: "mostly_iq4_nl",
    26: "mostly_iq3_s", 27: "mostly_iq3_m", 28: "mostly_iq2_s",
    29: "mostly_iq2_m", 30: "mostly_iq4_xs", 31: "mostly_iq1_m",
    32: "mostly_bf16", 36: "mostly_tq1_0", 37: "mostly_tq2_0",
    38: "mostly_mxfp4_moe", 39: "mostly_nvfp4", 40: "mostly_q1_0",
    41: "mostly_q2_0", 1024: "guessed",
}

HEX40 = re.compile(r"^[0-9a-f]{40}$")
# mirror common/download.cpp:get_gguf_split_info()
SPLIT_PART_RE = re.compile(r"^(.+)-(\d{5})-of-(\d{5})$", re.IGNORECASE)
TAG_RE = re.compile(r"[-.]([A-Za-z0-9_]+)$")


class NeedMore(Exception):
    """Raised by the parser when the buffer must grow."""


# ---- byte sources --------------------------------------------------------

class RemoteBuf:
    """Lazy byte source over HTTP with Range requests."""

    GUARD = 256 * MiB  # refuse to slurp absurdly large headers

    def __init__(self, url, token=None):
        self.url = url
        self.token = token
        self.buf = bytearray()

    def ensure(self, off, n):
        end = off + n
        while end > len(self.buf):
            if len(self.buf) >= self.GUARD:
                raise RuntimeError(f"GGUF header exceeded {self.GUARD // MiB} MiB; giving up")
            # fetch generously so a normal GGUF header lands in one request;
            # metadata + tensor table is usually well under a few MiB
            want = max(end - len(self.buf), 1 << 22)  # >= 4 MiB
            hdrs = {"Range": f"bytes={len(self.buf)}-{len(self.buf) + want - 1}"}
            if self.token:
                hdrs["Authorization"] = f"Bearer {self.token}"
            req = urllib.request.Request(self.url, headers=hdrs)
            with urllib.request.urlopen(req, timeout=30) as r:
                chunk = r.read(want)
            if not chunk:
                raise RuntimeError("server stopped sending header bytes early")
            self.buf += chunk
        return self.buf, off

    def content_length(self):
        hdrs = {}
        if self.token:
            hdrs["Authorization"] = f"Bearer {self.token}"
        req = urllib.request.Request(self.url, method="HEAD", headers=hdrs)
        with urllib.request.urlopen(req, timeout=30) as r:
            cl = r.headers.get("Content-Length")
            return int(cl) if cl else None

    def close(self):
        pass


class LocalBuf:
    """Lazy byte source over a local file. Reads only the header, not the
    whole file, so it is safe on multi-GB GGUFs."""

    def __init__(self, path):
        self.path = str(path)
        self.size = os.path.getsize(self.path)
        self.f = open(self.path, "rb")
        self.buf = bytearray()

    def ensure(self, off, n):
        end = off + n
        if end > self.size:
            raise RuntimeError(f"GGUF header read past EOF on {self.path} (truncated?)")
        if end > len(self.buf):
            want_end = min(max(end, len(self.buf) + (1 << 22)), self.size)
            self.f.seek(len(self.buf))
            data = self.f.read(want_end - len(self.buf))
            if not data:
                raise RuntimeError(f"unexpected EOF reading {self.path}")
            self.buf += data
        return self.buf, off

    def content_length(self):
        return self.size

    def close(self):
        try:
            self.f.close()
        except Exception:
            pass


# ---- low-level readers (every read bounds-checks and raises NeedMore) -----

def chk(buf, o, n):
    if o + n > len(buf):
        raise NeedMore(o + n)
    return buf, o


def rd_u8(buf, o):    buf, o = chk(buf, o, 1); return buf[o], o + 1
def rd_bool(buf, o):  buf, o = chk(buf, o, 1); return buf[o] != 0, o + 1
def rd_u16(buf, o):   buf, o = chk(buf, o, 2); return struct.unpack_from("<H", buf, o)[0], o + 2
def rd_i16(buf, o):   buf, o = chk(buf, o, 2); return struct.unpack_from("<h", buf, o)[0], o + 2
def rd_u32(buf, o):   buf, o = chk(buf, o, 4); return struct.unpack_from("<I", buf, o)[0], o + 4
def rd_i32(buf, o):   buf, o = chk(buf, o, 4); return struct.unpack_from("<i", buf, o)[0], o + 4
def rd_u64(buf, o):   buf, o = chk(buf, o, 8); return struct.unpack_from("<Q", buf, o)[0], o + 8
def rd_i64(buf, o):   buf, o = chk(buf, o, 8); return struct.unpack_from("<q", buf, o)[0], o + 8
def rd_f32(buf, o):   buf, o = chk(buf, o, 4); return struct.unpack_from("<f", buf, o)[0], o + 4
def rd_f64(buf, o):   buf, o = chk(buf, o, 8); return struct.unpack_from("<d", buf, o)[0], o + 8


def rd_str(buf, o):
    n, o = rd_u64(buf, o)
    buf, o = chk(buf, o, n)
    s = bytes(buf[o:o + n]).decode("utf-8", "replace")
    return s, o + n


def rd_value(buf, o, vtype):
    if vtype == VT_UINT8:   return rd_u8(buf, o)
    if vtype == VT_INT8:    buf, o = chk(buf, o, 1); return struct.unpack_from("<b", buf, o)[0], o + 1
    if vtype == VT_UINT16:  return rd_u16(buf, o)
    if vtype == VT_INT16:   return rd_i16(buf, o)
    if vtype == VT_UINT32:  return rd_u32(buf, o)
    if vtype == VT_INT32:   return rd_i32(buf, o)
    if vtype == VT_FLOAT32: return rd_f32(buf, o)
    if vtype == VT_BOOL:    return rd_bool(buf, o)
    if vtype == VT_STRING:  return rd_str(buf, o)
    if vtype == VT_UINT64:  return rd_u64(buf, o)
    if vtype == VT_INT64:   return rd_i64(buf, o)
    if vtype == VT_FLOAT64: return rd_f64(buf, o)
    if vtype == VT_ARRAY:
        sub, o = rd_u32(buf, o)
        cnt, o = rd_u64(buf, o)
        vals = []
        for _ in range(cnt):
            v, o = rd_value(buf, o, sub)
            vals.append(v)
        return vals, o
    raise ValueError(f"unknown value type {vtype}")


# ---- GGUF header parser ---------------------------------------------------

def parse_gguf_header(src):
    """Return (fields, tensors, data_offset). Raises NeedMore to grow."""
    buf, o = src.ensure(0, 12)
    magic, o = rd_u32(buf, o)
    if magic != GGUF_MAGIC:
        raise ValueError(f"not a GGUF file (magic=0x{magic:08x})")
    version, o = rd_u32(buf, o)
    if version != 3:
        sys.stderr.write(f"warning: unexpected GGUF version {version}\n")
    tensor_count, o = rd_u64(buf, o)
    kv_count, o = rd_u64(buf, o)

    fields = {}
    for _ in range(kv_count):
        name, o = rd_str(buf, o)
        vt, o = rd_u32(buf, o)
        val, o = rd_value(buf, o, vt)
        fields[name] = val

    tensors = []
    for _ in range(tensor_count):
        name, o = rd_str(buf, o)
        ndim, o = rd_u32(buf, o)
        dims = []
        for _ in range(ndim):
            d, o = rd_u64(buf, o)
            dims.append(d)
        ttype, o = rd_u32(buf, o)
        offset, o = rd_u64(buf, o)
        tensors.append({"name": name, "ndim": ndim, "dims": dims, "type": ttype, "offset": offset})

    align = DEFAULT_ALIGN
    al = fields.get("general.alignment")
    if isinstance(al, int):
        align = al
    pad = (-o) % align
    o += pad
    return fields, tensors, o


def parse_with_retry(src):
    while True:
        try:
            return parse_gguf_header(src)
        except NeedMore as e:
            need = e.args[0] if e.args else len(getattr(src, "buf", b"")) + 1
            src.ensure(need, 1)


# ---- HF Hub cache (mirrors common/hf-cache.cpp) ---------------------------

def hf_cache_root():
    """Resolve the HF Hub cache dir using llama.cpp's precedence chain."""
    for var in ("LLAMA_CACHE", "HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"):
        v = os.environ.get(var)
        if v:
            return pathlib.Path(v)
    v = os.environ.get("HF_HOME")
    if v:
        return pathlib.Path(v) / "hub"
    v = os.environ.get("XDG_CACHE_HOME")
    if v:
        return pathlib.Path(v) / "huggingface" / "hub"
    return pathlib.Path.home() / ".cache" / "huggingface" / "hub"


def hf_repo_dir(repo):
    return hf_cache_root() / ("models--" + repo.replace("/", "--"))


def hf_cached_commit(repo_dir):
    """Prefer refs/main, else first valid ref. Mirrors get_cached_ref()."""
    refs = repo_dir / "refs"
    if not refs.is_dir():
        return None
    fallback = None
    for rf in sorted(refs.iterdir()):
        if not rf.is_file():
            continue
        try:
            commit = rf.read_text(errors="replace").strip()
        except OSError:
            continue
        if not HEX40.match(commit):
            continue
        if rf.name == "main":
            return commit
        if fallback is None:
            fallback = commit
    return fallback


def hf_cached_ggufs(repo, fname=None):
    """Return list of (relpath, abspath) for cached .gguf files, or [] if the
    repo is not cached. If fname is given, restrict to that file."""
    rd = hf_repo_dir(repo)
    snaps = rd / "snapshots"
    if not snaps.is_dir():
        return []
    commit = hf_cached_commit(rd)
    commits = []
    if commit:
        commits.append(commit)
    for d in sorted(snaps.iterdir()):
        if d.is_dir() and d.name not in commits:
            commits.append(d.name)

    out = []
    seen = set()
    for c in commits:
        cdir = snaps / c
        if not cdir.is_dir():
            continue
        found_here = []
        for p in cdir.rglob("*"):
            if not p.is_file():  # follows symlinks; skips dangling ones
                continue
            rel = p.relative_to(cdir).as_posix()
            if not rel.lower().endswith(".gguf"):
                continue
            if fname and os.path.basename(rel) != fname and rel != fname:
                continue
            if rel in seen:
                continue
            seen.add(rel)
            found_here.append((rel, str(p)))
        if found_here:
            out.extend(found_here)
            break  # use the preferred commit only
    return out


# ---- HF online plumbing ---------------------------------------------------

def hf_headers(token):
    return {"Authorization": f"Bearer {token}"} if token else {}


def list_repo_files(repo, token, revision="main"):
    url = f"https://huggingface.co/api/models/{repo}"
    req = urllib.request.Request(url, headers=hf_headers(token))
    with urllib.request.urlopen(req, timeout=30) as r:
        meta = json.load(r)
    return [s["rfilename"] for s in meta.get("siblings", [])]


def resolve_url(repo, fname, revision="main"):
    return f"https://huggingface.co/{repo}/resolve/{revision}/{fname}"


def split_repo_tag(s):
    """Mirror common_download_split_repo_tag(): 'org/model:Q6_K' -> ('org/model', 'Q6_K')."""
    parts = s.split(":")
    tag = parts[-1] if len(parts) > 1 else ""
    repo = parts[0]
    return repo, tag


def gguf_split_info(name):
    """Mirror common/download.cpp:get_gguf_split_info(). Returns (prefix, tag, index, count).
    prefix retains the tag token; split parts share a prefix."""
    prefix = os.path.basename(name)
    if not prefix.lower().endswith(".gguf"):
        return "", "", 1, 1  # not a gguf; mirrors get_gguf_split_info early-return
    prefix = prefix[:-5]
    index, count = 1, 1
    m = SPLIT_PART_RE.match(prefix)
    if m:
        prefix = m.group(1)
        index = int(m.group(2))
        count = int(m.group(3))
    tag = ""
    mt = TAG_RE.search(prefix)
    if mt:
        tag = mt.group(1).upper()
    return prefix, tag, index, count


def build_groups(pairs):
    """pairs: [(kind, loc), ...] -> [{label, tag, parts:[(kind,loc),...]}, ...].
    Groups multi-part splits by their shared prefix; parts sorted by split index."""
    buckets = {}
    for kind, loc in pairs:
        prefix, tag, index, count = gguf_split_info(loc)
        b = buckets.setdefault(prefix, {"tag": "", "parts": []})
        b["parts"].append((index, kind, loc))
        if tag:
            b["tag"] = tag
    out = []
    for prefix, b in buckets.items():
        parts = sorted(b["parts"], key=lambda x: x[0])
        out.append({"label": prefix, "tag": b["tag"],
                    "parts": [(kind, loc) for _, kind, loc in parts]})
    return out


def filter_by_tag(groups, tag):
    if not tag:
        return groups
    tu = tag.upper()
    return [g for g in groups if g["tag"].upper() == tu]


def available_tags(groups):
    return sorted({g["tag"] for g in groups if g["tag"]})


def make_buf(desc, token):
    kind, loc = desc
    return LocalBuf(loc) if kind == "local" else RemoteBuf(loc, token)


# ---- math -----------------------------------------------------------------

def bpp_of(dt):
    if dt not in CACHE_BPP:
        raise SystemExit(f"unsupported cache dtype '{dt}'. one of: {sorted(CACHE_BPP)}")
    return CACHE_BPP[dt]


def weight_bytes(tensors, data_offset, content_length):
    """Sum tensor extents within ONE part via consecutive offsets."""
    if not tensors:
        return 0, {}
    ordered = sorted(tensors, key=lambda t: t["offset"])
    total = 0
    by_type = {}
    for i, t in enumerate(ordered):
        nxt = ordered[i + 1]["offset"] if i + 1 < len(ordered) \
            else (content_length - data_offset)
        sz = nxt - t["offset"]
        total += sz
        by_type[t["type"]] = by_type.get(t["type"], 0) + sz
    return total, by_type


def kv_bytes(hp, ctx, ct_k, ct_v, seqs):
    bk = bpp_of(ct_k) / 8
    bv = bpp_of(ct_v) / 8
    per_tok = hp["n_embd_k_gqa"] * bk + hp["n_embd_v_gqa"] * bv
    return per_tok * ctx * hp["n_layer"] * seqs


def human(n):
    if n >= GiB:
        return f"{n / GiB:.2f} GiB"
    if n >= MiB:
        return f"{n / MiB:.1f} MiB"
    return f"{n / 1024:.0f} KiB"


# ---- hparams extraction ----------------------------------------------------

def extract_hparams(fields):
    arch = fields.get("general.architecture", "")
    if not arch:
        raise SystemExit("no general.architecture in metadata; cannot infer hparams")
    g = lambda k: fields.get(f"{arch}.{k}")
    n_layer = g("block_count")
    n_embd = g("embedding_length")
    n_head = g("attention.head_count")
    n_head_kv = g("attention.head_count_kv") or n_head
    head_k = g("attention.key_length") or g("rope.dimension_count")
    head_v = g("attention.value_length") or head_k
    if head_k is None and n_head:
        head_k = n_embd // n_head
    if head_v is None:
        head_v = head_k
    return {
        "arch": arch,
        "n_layer": n_layer,
        "n_embd": n_embd,
        "n_head": n_head,
        "n_head_kv": n_head_kv,
        "n_embd_head_k": head_k,
        "n_embd_head_v": head_v,
        "n_embd_k_gqa": (head_k or 0) * (n_head_kv or 0),
        "n_embd_v_gqa": (head_v or 0) * (n_head_kv or 0),
        "n_ctx_train": g("context_length"),
        "name": fields.get("general.name"),
        "ftype": fields.get("general.file_type"),
    }


# ---- group analysis -------------------------------------------------------

def estimatable(hp):
    """Return '' if we can estimate memory for this model, else a short reason.
    Guards out mmproj/clip/encoder GGUFs that ship alongside LLMs in a repo."""
    if hp["n_layer"] is None:
        return "no block_count (not a transformer LLM, e.g. mmproj/clip)"
    if not hp["n_embd_k_gqa"] or not hp["n_embd_v_gqa"]:
        return "cannot infer attention head dims (no head_count/rope.dim)"
    return ""


def compute_group(part_descs, args):
    """Parse every part of a (possibly split) model. Weights are summed
    per-part (each part has its own header/tensors, per llama-model-loader.cpp),
    so split accounting is correct. Hparams come from whichever part carries
    general.architecture (part 1 in practice)."""
    hp = None
    total_wb = 0
    total_bt = {}
    total_clen = 0
    for desc in part_descs:
        src = make_buf(desc, args.token)
        try:
            fields, tensors, data_off = parse_with_retry(src)
            if hp is None and "general.architecture" in fields:
                hp = extract_hparams(fields)
            plen = src.content_length()
            total_clen += plen
            wb, bt = weight_bytes(tensors, data_off, plen)
            total_wb += wb
            for t, b in bt.items():
                total_bt[t] = total_bt.get(t, 0) + b
        finally:
            src.close()
    if hp is None:
        raise SystemExit("no general.architecture found in any part")
    skip = estimatable(hp)
    kv = 0 if skip else kv_bytes(hp, args.ctx, args.ctk, args.ctv, args.seqs)
    return hp, total_wb, total_bt, kv, total_clen, skip


def show_detail(hp, wb, by_type, kv, ctx, ct_k, ct_v, seqs, ngl, content_length):
    print("=" * 64)
    nm = hp.get("name") or "(unnamed)"
    ft = FTYPE_NAMES.get(hp.get("ftype"), f"#{hp.get('ftype')}")
    print(f"  model      : {nm}")
    print(f"  arch       : {hp['arch']}   file_type: {ft}")
    print("-" * 64)
    print(f"  n_layer        : {hp['n_layer']}")
    print(f"  n_embd         : {hp['n_embd']}")
    print(f"  n_head / n_kv  : {hp['n_head']} / {hp['n_head_kv']}")
    print(f"  head dim k/v   : {hp['n_embd_head_k']} / {hp['n_embd_head_v']}")
    print(f"  n_embd_k_gqa   : {hp['n_embd_k_gqa']}")
    print(f"  n_ctx_train    : {hp['n_ctx_train']}")
    print(f"  gguf size      : {human(content_length)} ({content_length} B)")
    print("-" * 64)
    print(f"  weights (disk) : {human(wb)}")
    if by_type:
        for t, b in sorted(by_type.items(), key=lambda x: -x[1]):
            print(f"    {QUANT_NAMES.get(t, t):>10}: {human(b)}")
    print("-" * 64)
    print(f"  kv cache @{ctx} ctx, {ct_k}/{ct_v}, {seqs} seq(s): {human(kv)}")
    if ngl is not None and hp["n_layer"]:
        frac = min(ngl, hp["n_layer"]) / hp["n_layer"]
        print(f"    gpu offload {ngl}/{hp['n_layer']} layers -> ~{frac*100:.0f}% of KV on GPU")
    print("-" * 64)
    est = wb + kv
    print(f"  resident est  : {human(est)}  (weights + kv)")
    print(f"  + compute     : transient, ~few hundred MiB (varies with ubatch)")
    print("=" * 64)
    if hp["arch"] in RECURRENT_ARCHS:
        print("  NOTE: recurrent arch; KV-cache formula is for transformers.")
    if hp["n_ctx_train"] and ctx < hp["n_ctx_train"]:
        print(f"  NOTE: ctx {ctx} < train ctx {hp['n_ctx_train']}; "
              f"KV scales linearly (x{hp['n_ctx_train']/ctx:.1f} for full ctx).")


def run_detail(groups, args, note):
    if note:
        print(note)
    for g in groups:
        hp, wb, bt, kv, clen, skip = compute_group(g["parts"], args)
        print(f"\n>>> {g['label']}")
        if skip:
            print(f"  (skipped: {skip})")
            continue
        show_detail(hp, wb, bt, kv, args.ctx, args.ctk, args.ctv, args.seqs,
                    args.ngl, clen)


def run_summary(groups, args, note):
    print(f"\nRepo {args.model}  ({len(groups)} file{'s' if len(groups)!=1 else ''})  {note or ''}\n")
    print(f"{'file':<42}{'weights':>12}{'kv@'+str(args.ctx):>12}{'resident':>12}")
    print("-" * 78)
    rows = []
    skipped = []
    for g in groups:
        _, wb, _, kv, _, skip = compute_group(g["parts"], args)
        if skip:
            skipped.append(g["label"])
            continue
        rows.append((g["label"], wb, kv))
    for label, wb, kv in sorted(rows, key=lambda x: -x[1]):
        print(f"{label:<42}{human(wb):>12}{human(kv):>12}{human(wb+kv):>12}")
    if skipped:
        print(f"\nskipped (not LLMs): {', '.join(skipped)}")
    print("\n(re-run with --all or --file NAME for a full breakdown)\n")


def run_numbers(groups, args):
    """Non-verbose: print estimated resident memory (weights + kv) in MiB.
    One group -> a single integer. Several groups -> one integer per line,
    heaviest first (matches the summary-table order). Non-LLM GGUFs that
    ship in a repo (mmproj/clip) are silently skipped."""
    nums = []
    for g in groups:
        _, wb, _, kv, _, skip = compute_group(g["parts"], args)
        if skip:
            continue
        nums.append(round((wb + kv) / MiB))
    if not nums:
        sys.stderr.write("no estimatable LLM model found (only non-transformer GGUFs?)\n")
        sys.exit(1)
    if len(nums) == 1:
        print(nums[0])
    else:
        for n in sorted(nums, reverse=True):
            print(n)


def dispatch(groups, args, note):
    if not args.verbose:
        run_numbers(groups, args)
    elif args.all or len(groups) <= 1:
        run_detail(groups, args, note)
    else:
        run_summary(groups, args, note)


# ---- main -----------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("model", help=".gguf path or HF repo id, optionally org/model:quant")
    ap.add_argument("--file", help="specific .gguf file inside the HF repo")
    ap.add_argument("--ctx", type=int, default=4096, help="context size for KV estimate (default 4096)")
    ap.add_argument("-ctk", default="f16", help="KV key dtype (default f16)")
    ap.add_argument("-ctv", default="f16", help="KV value dtype (default f16)")
    ap.add_argument("--seqs", type=int, default=1, help="concurrent sequences (default 1)")
    ap.add_argument("--ngl", type=int, default=None, help="GPU layers offloaded (show KV split)")
    ap.add_argument("--token", default=os.environ.get("HF_TOKEN"))
    ap.add_argument("--revision", default="main")
    ap.add_argument("-v", "--verbose", action="store_true",
                    help="print the full breakdown (default: just estimated MiB)")
    ap.add_argument("--all", action="store_true", help="analyse every .gguf in the repo")
    ap.add_argument("--online", action="store_true",
                    help="ignore the local HF cache and fetch headers over HTTP")
    args = ap.parse_args()

    if os.path.isfile(args.model):
        dispatch([{"label": "", "tag": "", "parts": [("local", args.model)]}], args, None)
        return

    # HF repo id, optionally with a ':tag' quant selector (llama.cpp router style).
    repo, tag = split_repo_tag(args.model)
    if repo.count("/") != 1:
        raise SystemExit(
            f"'{args.model}' is neither an existing file nor a valid HF repo id "
            "(expected <org>/<model>[:quant])"
        )

    # Prefer the local HF cache (no re-fetch). Fall back to HTTP only if the
    # repo (or the requested tag) is absent locally.
    if not args.online:
        cached = hf_cached_ggufs(repo, fname=args.file)
        if cached:
            groups = build_groups([("local", p) for _, p in cached])
            if tag:
                groups = filter_by_tag(groups, tag)
            if groups:
                note = f"(from local cache: {hf_repo_dir(repo)})"
                dispatch(groups, args, note)
                return
            if tag:
                sys.stderr.write(
                    f"note: tag '{tag}' not found in local cache; "
                    "falling back to online\n"
                )
            # else: repo cached but somehow no ggufs -> fall through to online

    files = list_repo_files(repo, args.token, args.revision)
    ggufs = [f for f in files if f.lower().endswith(".gguf")]
    if not ggufs:
        raise SystemExit(f"no .gguf files found in repo {repo}")

    if args.file:
        sel = [f for f in ggufs if os.path.basename(f) == args.file or f == args.file]
        if not sel:
            raise SystemExit(f"--file {args.file} not found among: {ggufs}")
        ggufs = sel

    groups = build_groups([("remote", resolve_url(repo, f, args.revision)) for f in ggufs])
    if tag:
        groups = filter_by_tag(groups, tag)
        if not groups:
            avail = ", ".join(available_tags(build_groups(
                [("remote", u) for u in [resolve_url(repo, f, args.revision) for f in ggufs]])))
            raise SystemExit(
                f"tag '{tag}' not found in repo {repo}. "
                f"Available tags: {avail or '(none)'}"
            )
    dispatch(groups, args, "(from huggingface.co over HTTP)")


if __name__ == "__main__":
    main()
