#!/usr/bin/env python3
"""
split_gguf.py — Layer-aware GGUF splitter / byte-split / join / status

Commands:
  split_layers  — Split by transformer layers (primary, recommended)
                  Each output is a valid GGUF compatible with llama.cpp split loading.
                  llama.cpp / llama-cpp-python loads the first file — auto-discovers the rest.
  split         — Split by fixed byte size (legacy .part files)
  join          — Rejoin legacy .part files into one GGUF
  status        — Show all chunk / GGUF state

Usage:
  python split_gguf.py split_layers            # layer split (default 100MB target)
  python split_gguf.py split_layers --mb 50    # smaller groups
  python split_gguf.py split                   # legacy byte split
  python split_gguf.py join                    # rejoin legacy .part files
  python split_gguf.py status                  # show all state
"""

import os, sys, json, argparse, struct

# ── GGUF constants ─────────────────────────────────────────────────────────────
GGUF_MAGIC = b'GGUF'
GGUF_ALIGN = 32

# KV value types
U8=0; I8=1; U16=2; I16=3; U32=4; I32=5; F32=6; BOOL=7; STRING=8; ARRAY=9
U64=10; I64=11; F64=12

TYPE_SIZES = {U8:1,I8:1,U16:2,I16:2,U32:4,I32:4,F32:4,BOOL:1,U64:8,I64:8,F64:8}


def _align(n, a=GGUF_ALIGN):
    return ((n + a - 1) // a) * a


def _load_cfg():
    p = os.path.join(os.path.dirname(os.path.abspath(__file__)), "install.json")
    try:
        with open(p) as f:
            c = json.load(f)
        g = c.get("gguf", {})
        return {
            "dir":      g.get("dir",      "./model_gguf"),
            "default":  g.get("default",  "Qwen3-1.7B-Q4_K_M.gguf"),
            "quality":  g.get("quality",  "Qwen3-1.7B-Q8_0.gguf"),
            "chunk_mb": int(c.get("chunks", {}).get("chunk_mb", 100)),
        }
    except Exception:
        return {"dir":"./model_gguf","default":"Qwen3-1.7B-Q4_K_M.gguf",
                "quality":"Qwen3-1.7B-Q8_0.gguf","chunk_mb":100}


# ══════════════════════════════════════════════════════════════════════════════
#  GGUF Parser
# ══════════════════════════════════════════════════════════════════════════════

def _read_str(f):
    n = int.from_bytes(f.read(8), 'little')
    return f.read(n).decode('utf-8', errors='replace')


def _skip_value(f, vtype):
    if vtype == STRING:
        f.read(int.from_bytes(f.read(8), 'little'))
    elif vtype == ARRAY:
        atype = int.from_bytes(f.read(4), 'little')
        alen  = int.from_bytes(f.read(8), 'little')
        if atype == STRING:
            for _ in range(alen):
                f.read(int.from_bytes(f.read(8), 'little'))
        else:
            f.read(alen * TYPE_SIZES.get(atype, 4))
    else:
        f.read(TYPE_SIZES.get(vtype, 4))


def _read_kv_entries(path, n_kv, kv_start):
    """
    Read KV entries one-by-one from a GGUF file.
    Returns list of (key_str, raw_bytes) tuples.
    raw_bytes = complete binary representation of that entry.
    """
    entries = []
    with open(path, 'rb') as f:
        f.seek(kv_start)
        for _ in range(n_kv):
            entry_start = f.tell()
            key_len = int.from_bytes(f.read(8), 'little')
            key     = f.read(key_len).decode('utf-8', errors='replace')
            vtype   = int.from_bytes(f.read(4), 'little')
            _skip_value(f, vtype)
            entry_end = f.tell()
            f.seek(entry_start)
            raw = f.read(entry_end - entry_start)
            entries.append((key, raw))
    return entries


def _parse_gguf(path):
    """
    Parse a GGUF file.
    Returns dict: version, n_tensors, n_kv, kv_raw (bytes), tensors (list), data_start.
    Each tensor: {name, shape, dtype, offset, size}
    offset = byte offset from data_start (in source file).
    size   = byte size of tensor data.
    """
    file_size = os.path.getsize(path)
    with open(path, 'rb') as f:
        magic = f.read(4)
        if magic != GGUF_MAGIC:
            raise ValueError(f"Not a GGUF file (magic={magic!r})")
        version   = int.from_bytes(f.read(4), 'little')
        n_tensors = int.from_bytes(f.read(8), 'little')
        n_kv      = int.from_bytes(f.read(8), 'little')

        # Capture raw KV bytes (copied verbatim to each output split)
        kv_start = f.tell()
        for _ in range(n_kv):
            f.read(int.from_bytes(f.read(8), 'little'))   # key
            _skip_value(f, int.from_bytes(f.read(4), 'little'))
        kv_end = f.tell()
        f.seek(kv_start)
        kv_raw = f.read(kv_end - kv_start)

        # Tensor infos — preserve original declaration order
        tensors = []
        for _ in range(n_tensors):
            name   = _read_str(f)
            n_dims = int.from_bytes(f.read(4), 'little')
            shape  = [int.from_bytes(f.read(8), 'little') for _ in range(n_dims)]
            dtype  = int.from_bytes(f.read(4), 'little')
            offset = int.from_bytes(f.read(8), 'little')
            tensors.append({'name': name, 'shape': shape,
                            'dtype': dtype, 'offset': offset})

        # Data section starts at next 32-byte boundary after tensor infos
        data_start = _align(f.tell())

    # Compute tensor data sizes from consecutive offsets (sort by offset)
    by_offset = sorted(tensors, key=lambda t: t['offset'])
    total_data = file_size - data_start
    for i, t in enumerate(by_offset):
        nxt = by_offset[i+1]['offset'] if i+1 < len(by_offset) else total_data
        t['size'] = nxt - t['offset']

    return {
        '_path':      path,
        'version':    version,
        'n_tensors':  n_tensors,
        'n_kv':       n_kv,
        'kv_raw':     kv_raw,
        'tensors':    tensors,
        'data_start': data_start,
        'file_size':  file_size,
    }


# ══════════════════════════════════════════════════════════════════════════════
#  GGUF Writer (one split)
# ══════════════════════════════════════════════════════════════════════════════

def _kv_entry(key, vtype, value):
    """Serialize one KV pair."""
    k = key.encode('utf-8')
    b = len(k).to_bytes(8, 'little') + k + vtype.to_bytes(4, 'little')
    if   vtype == U16:    b += struct.pack('<H', value)
    elif vtype == U32:    b += struct.pack('<I', value)
    elif vtype == STRING:
        v = value.encode('utf-8')
        b += len(v).to_bytes(8, 'little') + v
    return b


def _write_split(out_path, gguf, group_tensors, split_no, split_count):
    """
    Write one layer-split GGUF file.

    group_tensors : tensor dicts for this split (subset of gguf['tensors'])
    split_no      : 0-based index of this split
    split_count   : total number of splits

    Adds split.no, split.count, split.tensors_count KV entries so llama.cpp
    can auto-load all splits from the first file.
    """
    src_path      = gguf['_path']
    version       = gguf['version']
    kv_raw        = gguf['kv_raw']
    n_kv_orig     = gguf['n_kv']
    data_start    = gguf['data_start']
    total_tensors = gguf['n_tensors']

    # Split metadata KV entries (appended after original KV)
    split_kv  = _kv_entry('split.no',             U16, split_no)
    split_kv += _kv_entry('split.count',           U16, split_count)
    split_kv += _kv_entry('split.tensors_count',   U32, total_tensors)
    n_kv_total = n_kv_orig + 3

    # Sort group by original data offset for sequential reading
    grp_sorted = sorted(group_tensors, key=lambda t: t['offset'])

    # Assign new data offsets within this split's data section (32-byte aligned)
    new_off = 0
    for t in grp_sorted:
        t['new_offset'] = new_off
        new_off = _align(new_off + t['size'])

    # Build tensor info bytes (preserve original declaration order for this group)
    ti_bytes = b''
    for t in group_tensors:
        name_enc = t['name'].encode('utf-8')
        ti_bytes += len(name_enc).to_bytes(8, 'little') + name_enc
        ti_bytes += len(t['shape']).to_bytes(4, 'little')
        for dim in t['shape']:
            ti_bytes += dim.to_bytes(8, 'little')
        ti_bytes += t['dtype'].to_bytes(4, 'little')
        ti_bytes += t['new_offset'].to_bytes(8, 'little')

    with open(out_path, 'wb') as out, open(src_path, 'rb') as src:
        # ── Header ──
        out.write(GGUF_MAGIC)
        out.write(version.to_bytes(4, 'little'))
        out.write(len(group_tensors).to_bytes(8, 'little'))
        out.write(n_kv_total.to_bytes(8, 'little'))

        # ── KV section ──
        out.write(kv_raw)      # original metadata
        out.write(split_kv)    # split.no / split.count / split.tensors_count

        # ── Tensor infos ──
        out.write(ti_bytes)

        # ── Align to data section ──
        cur = out.tell()
        pad = _align(cur) - cur
        if pad:
            out.write(b'\x00' * pad)

        # ── Tensor data (in new_offset order = grp_sorted order) ──
        write_pos = 0
        for t in grp_sorted:
            # Pad to this tensor's aligned start
            pad = t['new_offset'] - write_pos
            if pad > 0:
                out.write(b'\x00' * pad)
            # Copy tensor data from source
            src.seek(data_start + t['offset'])
            remaining = t['size']
            buf_size  = 4 * 1024 * 1024
            while remaining > 0:
                buf = src.read(min(buf_size, remaining))
                if not buf:
                    break
                out.write(buf)
                remaining -= len(buf)
            write_pos = t['new_offset'] + t['size']


# ══════════════════════════════════════════════════════════════════════════════
#  Layer grouping
# ══════════════════════════════════════════════════════════════════════════════

def _group_by_layers(tensors, target_mb):
    """
    Classify tensors by transformer layer, then merge adjacent layers
    until each group reaches target_mb.

    Layer classes:
      embed   — token_embd.*, token_embd_norm.*
      blk.N   — blk.N.*  (transformer blocks 0..N)
      output  — output.*, output_norm.*

    Returns: list of (label_str, [tensor, ...])
    """
    # ── Classify ──
    buckets = {}  # (category, sort_key) → [tensors]
    for t in tensors:
        name = t['name']
        if name.startswith('blk.'):
            n = int(name.split('.')[1])
            key = ('blk', n)
        elif name.startswith('token_embd'):
            key = ('embed', -1)
        else:
            key = ('output', 999999)
        buckets.setdefault(key, []).append(t)

    # Sort: embed first, then blk.0…blk.N in order, output last
    sorted_buckets = sorted(buckets.items(),
        key=lambda x: (0 if x[0][0]=='embed' else
                       1 if x[0][0]=='blk'   else 2,  x[0][1]))

    # ── Compute MB per bucket ──
    bucket_info = []
    for (cat, n), tlist in sorted_buckets:
        mb = sum(t['size'] for t in tlist) / (1024**2)
        if cat == 'embed':   label = 'embed'
        elif cat == 'blk':   label = f'blk{n:02d}'
        else:                label = 'output'
        bucket_info.append((label, tlist, mb))

    # ── Merge buckets greedily to reach target_mb per group ──
    target = max(float(target_mb), 1.0)
    groups = []
    cur_t, cur_mb, cur_label = [], 0.0, ''

    for label, tlist, mb in bucket_info:
        # Start new group if adding this bucket would exceed target (and we have something)
        if cur_t and cur_mb + mb > target:
            groups.append((cur_label, cur_t))
            cur_t, cur_mb, cur_label = [], 0.0, ''
        cur_t.extend(tlist)
        cur_mb += mb
        cur_label = f"{cur_label}+{label}" if cur_label else label

    if cur_t:
        groups.append((cur_label, cur_t))

    return groups


# ══════════════════════════════════════════════════════════════════════════════
#  Public split_layers
# ══════════════════════════════════════════════════════════════════════════════

def split_gguf_layers(model_path, target_mb=100, out_dir=None):
    """
    Layer-aware GGUF split.

    Each output file:
      - Is a valid standalone GGUF (can be inspected with gguf-dump etc.)
      - Contains split.no / split.count / split.tensors_count metadata
      - Is named <base>-00001-of-NNNNN.gguf

    llama.cpp / llama-cpp-python auto-loads all splits when given the first file:
      Llama('./model_gguf/chunks/Model-00001-of-00010.gguf')

    Returns list of created output file paths.
    """
    if not os.path.isfile(model_path):
        print(f"[split_layers] ERROR: File not found: {model_path}")
        return []

    base     = os.path.splitext(os.path.basename(model_path))[0]
    gguf_dir = os.path.dirname(model_path)
    if out_dir is None:
        out_dir = os.path.join(gguf_dir, "chunks")
    os.makedirs(out_dir, exist_ok=True)

    file_mb = os.path.getsize(model_path) / (1024**2)
    print(f"[split_layers] Parsing {os.path.basename(model_path)} ({file_mb:.0f}MB) ...")
    gguf   = _parse_gguf(model_path)
    groups = _group_by_layers(gguf['tensors'], target_mb)
    n      = len(groups)

    print(f"[split_layers] {gguf['n_tensors']} tensors  →  "
          f"{n} groups  (target {target_mb}MB/group)")
    print(f"[split_layers] Output: {out_dir}/")
    print()

    parts = []
    for i, (label, tensors_grp) in enumerate(groups):
        out_name = f"{base}-{i+1:05d}-of-{n:05d}.gguf"
        out_path = os.path.join(out_dir, out_name)
        grp_mb   = sum(t['size'] for t in tensors_grp) / (1024**2)
        print(f"  [{i+1:2d}/{n}]  {out_name}  "
              f"({len(tensors_grp):3d} tensors, {grp_mb:5.1f}MB)  [{label}]",
              flush=True)
        _write_split(out_path, gguf, tensors_grp, split_no=i, split_count=n)
        parts.append(out_path)

    total_out_mb = sum(os.path.getsize(p) for p in parts) / (1024**2)
    print()
    print(f"[split_layers] Done!  {n} files  ({total_out_mb:.0f}MB total)  → {out_dir}/")
    print(f"[split_layers] Load:  Llama('{parts[0]}')")
    return parts


# ══════════════════════════════════════════════════════════════════════════════
#  Legacy byte-split (kept for backward compat)
# ══════════════════════════════════════════════════════════════════════════════

def split_gguf(model_path, chunk_mb=100):
    """Split a GGUF file into raw <chunk_mb>MB .part files (legacy)."""
    if not os.path.isfile(model_path):
        print(f"[split] ERROR: File not found: {model_path}")
        return []

    base      = os.path.splitext(os.path.basename(model_path))[0]
    gguf_dir  = os.path.dirname(model_path)
    out_dir   = os.path.join(gguf_dir, "chunks")
    os.makedirs(out_dir, exist_ok=True)

    chunk_bytes = chunk_mb * 1024 * 1024
    file_size   = os.path.getsize(model_path)
    n_chunks    = (file_size + chunk_bytes - 1) // chunk_bytes

    print(f"[split] {os.path.basename(model_path)} ({file_size//(1024**2)}MB)"
          f"  →  {n_chunks} × {chunk_mb}MB chunks")
    print(f"[split] Output: {out_dir}/")
    print()

    parts = []
    with open(model_path, 'rb') as src:
        for i in range(1, n_chunks + 1):
            part_name = f"{base}-{i:05d}-of-{n_chunks:05d}.part"
            part_path = os.path.join(out_dir, part_name)
            written   = 0
            with open(part_path, 'wb') as dst:
                while written < chunk_bytes:
                    buf = src.read(min(4 * 1024 * 1024, chunk_bytes - written))
                    if not buf:
                        break
                    dst.write(buf)
                    written += len(buf)
            parts.append(part_path)
            pct = int(i / n_chunks * 40)
            bar = '█' * pct + '░' * (40 - pct)
            print(f"\r  [{bar}] {i}/{n_chunks}  {part_name}", end='', flush=True)

    print(f"\n\n[split] Done! {n_chunks} chunks → {out_dir}/")
    print(f"[split] Rejoin: cat {out_dir}/{base}-*.part > {model_path}")
    return parts


# ══════════════════════════════════════════════════════════════════════════════
#  Join layer-split GGUF files → single combined GGUF
# ══════════════════════════════════════════════════════════════════════════════

def join_gguf_layers(split_paths, out_path):
    """
    Merge layer-split GGUF files (e.g. Model-00001-of-00011.gguf, ...) into
    one combined GGUF at out_path.

    Steps:
      1. Parse each split file to get tensor metadata + data layout.
      2. Read KV entries from the first split, filter out split.* entries.
      3. Collect all tensors from all splits in declaration order.
      4. Assign new 32-byte-aligned offsets for the combined data section.
      5. Write header → filtered KV → tensor infos → tensor data.

    Returns out_path on success, raises on failure.
    """
    SPLIT_KEYS = {'split.no', 'split.count', 'split.tensors_count'}

    print(f"[join_layers] Parsing {len(split_paths)} splits ...", flush=True)
    splits = [_parse_gguf(p) for p in split_paths]
    first  = splits[0]

    # ── Filtered KV (remove split.* entries from first split) ──────────────
    # KV section starts at byte 24 (4 magic + 4 version + 8 n_tensors + 8 n_kv)
    kv_start = 24
    raw_entries    = _read_kv_entries(first['_path'], first['n_kv'], kv_start)
    filtered       = [(k, raw) for k, raw in raw_entries if k not in SPLIT_KEYS]
    filtered_bytes = b''.join(raw for _, raw in filtered)
    n_kv_filtered  = len(filtered)

    # ── Collect all tensors in split order, then by original offset ─────────
    all_tensors = []
    for split in splits:
        for t in sorted(split['tensors'], key=lambda x: x['offset']):
            all_tensors.append({**t, '_src': split})

    total = len(all_tensors)

    # ── Assign new 32-byte-aligned offsets ──────────────────────────────────
    new_off = 0
    for t in all_tensors:
        t['new_offset'] = new_off
        new_off = _align(new_off + t['size'])

    # ── Build tensor info bytes ──────────────────────────────────────────────
    ti_bytes = b''
    for t in all_tensors:
        name_enc  = t['name'].encode('utf-8')
        ti_bytes += len(name_enc).to_bytes(8, 'little') + name_enc
        ti_bytes += len(t['shape']).to_bytes(4, 'little')
        for dim in t['shape']:
            ti_bytes += dim.to_bytes(8, 'little')
        ti_bytes += t['dtype'].to_bytes(4, 'little')
        ti_bytes += t['new_offset'].to_bytes(8, 'little')

    out_mb = sum(t['size'] for t in all_tensors) / (1024**2)
    print(f"[join_layers] {total} tensors  ({out_mb:.0f}MB data)  "
          f"→  {os.path.basename(out_path)}", flush=True)

    src_handles: dict = {}
    try:
        with open(out_path, 'wb') as out:
            # ── GGUF Header ──
            out.write(GGUF_MAGIC)
            out.write(first['version'].to_bytes(4, 'little'))
            out.write(total.to_bytes(8, 'little'))
            out.write(n_kv_filtered.to_bytes(8, 'little'))

            # ── Filtered KV section ──
            out.write(filtered_bytes)

            # ── Tensor infos ──
            out.write(ti_bytes)

            # ── Pad to 32-byte data boundary ──
            cur = out.tell()
            pad = _align(cur) - cur
            if pad:
                out.write(b'\x00' * pad)

            # ── Tensor data ──
            write_pos = 0
            for i, t in enumerate(all_tensors, 1):
                src_path   = t['_src']['_path']
                data_start = t['_src']['data_start']

                # Pad to aligned offset
                gap = t['new_offset'] - write_pos
                if gap > 0:
                    out.write(b'\x00' * gap)

                # Open source file lazily
                if src_path not in src_handles:
                    src_handles[src_path] = open(src_path, 'rb')
                src = src_handles[src_path]
                src.seek(data_start + t['offset'])

                remaining = t['size']
                buf_size  = 4 * 1024 * 1024
                while remaining > 0:
                    buf = src.read(min(buf_size, remaining))
                    if not buf:
                        break
                    out.write(buf)
                    remaining -= len(buf)

                write_pos = t['new_offset'] + t['size']

                if i % 50 == 0 or i == total:
                    pct = int(i / total * 40)
                    bar = '█' * pct + '░' * (40 - pct)
                    print(f"\r  [{bar}] {i}/{total}", end='', flush=True)
    finally:
        for h in src_handles.values():
            try:
                h.close()
            except Exception:
                pass

    size_mb = os.path.getsize(out_path) / (1024**2)
    print(f"\n[join_layers] Done!  {size_mb:.0f}MB  →  {out_path}", flush=True)
    return out_path


# ══════════════════════════════════════════════════════════════════════════════
#  Join (legacy .part files only)
# ══════════════════════════════════════════════════════════════════════════════

def join_gguf(model_path):
    """Rejoin legacy .part chunks into the original GGUF file."""
    base      = os.path.splitext(os.path.basename(model_path))[0]
    gguf_dir  = os.path.dirname(model_path)
    chunk_dir = os.path.join(gguf_dir, "chunks")

    if not os.path.isdir(chunk_dir):
        print(f"[join] No chunks/ directory: {chunk_dir}")
        return False

    parts = sorted(f for f in os.listdir(chunk_dir)
                   if f.startswith(base) and f.endswith('.part'))
    if not parts:
        print(f"[join] No .part files for: {base}")
        return False

    if os.path.isfile(model_path):
        print(f"[join] Already exists: {model_path} (skipping)")
        return True

    total_mb = sum(os.path.getsize(os.path.join(chunk_dir, p)) for p in parts) / (1024**2)
    print(f"[join] Joining {len(parts)} parts → {os.path.basename(model_path)} ({total_mb:.0f}MB)")

    with open(model_path, 'wb') as out:
        for i, part in enumerate(parts, 1):
            with open(os.path.join(chunk_dir, part), 'rb') as f:
                while True:
                    buf = f.read(4 * 1024 * 1024)
                    if not buf:
                        break
                    out.write(buf)
            pct = int(i / len(parts) * 40)
            bar = '█' * pct + '░' * (40 - pct)
            print(f"\r  [{bar}] {i}/{len(parts)}  {part}", end='', flush=True)

    print(f"\n\n[join] Done! {os.path.basename(model_path)}  "
          f"({os.path.getsize(model_path)//(1024**2)}MB)")
    return True


# ══════════════════════════════════════════════════════════════════════════════
#  Status
# ══════════════════════════════════════════════════════════════════════════════

def status():
    cfg       = _load_cfg()
    gguf_dir  = cfg['dir']
    chunk_dir = os.path.join(gguf_dir, 'chunks')

    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("  GGUF Status")
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")

    for fname in [cfg['default'], cfg['quality']]:
        p = os.path.join(gguf_dir, fname)
        if os.path.isfile(p):
            mb = os.path.getsize(p) / (1024**2)
            print(f"  ✓  {fname}  ({mb:.0f}MB)")
        else:
            print(f"  ✗  {fname}  (not found)")

    print()
    if os.path.isdir(chunk_dir):
        # Layer-split .gguf files
        gguf_splits = sorted(f for f in os.listdir(chunk_dir) if f.endswith('.gguf'))
        part_chunks = sorted(f for f in os.listdir(chunk_dir) if f.endswith('.part'))

        if gguf_splits:
            total_mb = sum(os.path.getsize(os.path.join(chunk_dir, f))
                           for f in gguf_splits) / (1024**2)
            print(f"  Layer splits: {len(gguf_splits)} .gguf files  ({total_mb:.0f}MB total)")
            for f in gguf_splits:
                mb = os.path.getsize(os.path.join(chunk_dir, f)) / (1024**2)
                print(f"    {f}  ({mb:.1f}MB)")
        if part_chunks:
            total_mb = sum(os.path.getsize(os.path.join(chunk_dir, f))
                           for f in part_chunks) / (1024**2)
            print(f"  Byte chunks:  {len(part_chunks)} .part files  ({total_mb:.0f}MB total)")
        if not gguf_splits and not part_chunks:
            print("  Chunks: none")
        print(f"  Dir: {chunk_dir}/")
    else:
        print("  Chunks: none (no chunks/ directory)")
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")


# ══════════════════════════════════════════════════════════════════════════════
#  CLI
# ══════════════════════════════════════════════════════════════════════════════

def main():
    cfg          = _load_cfg()
    default_file = os.path.join(cfg['dir'], cfg['default'])

    p   = argparse.ArgumentParser(description="GGUF layer-aware splitter / joiner")
    sub = p.add_subparsers(dest='cmd')

    # split_layers
    sl = sub.add_parser('split_layers', help='Layer-aware split (recommended)')
    sl.add_argument('--file', default=default_file, help='GGUF file to split')
    sl.add_argument('--mb',   type=int, default=cfg['chunk_mb'],
                    help='Target MB per group (default from install.json)')

    # split (legacy)
    sp = sub.add_parser('split', help='Byte-size split (legacy .part files)')
    sp.add_argument('--file', default=default_file)
    sp.add_argument('--mb',   type=int, default=cfg['chunk_mb'])

    # join_layers (merge layer-split .gguf files)
    jl = sub.add_parser('join_layers', help='Merge layer-split .gguf files into one')
    jl.add_argument('--file', default=default_file,
                    help='Original model path (output destination)')
    jl.add_argument('--dir',  default=None,
                    help='chunks/ dir (default: <gguf_dir>/chunks)')

    # join (legacy .part files)
    jp = sub.add_parser('join', help='Rejoin legacy .part files')
    jp.add_argument('--file', default=default_file)

    # status
    sub.add_parser('status', help='Show GGUF and chunk state')

    args = p.parse_args()

    if   args.cmd == 'split_layers':
        split_gguf_layers(args.file, target_mb=args.mb)
    elif args.cmd == 'split':
        split_gguf(args.file, chunk_mb=args.mb)
    elif args.cmd == 'join_layers':
        base      = os.path.splitext(os.path.basename(args.file))[0]
        chunk_dir = args.dir or os.path.join(os.path.dirname(args.file), 'chunks')
        splits    = sorted(
            f for f in os.listdir(chunk_dir)
            if f.startswith(base) and f.endswith('.gguf')
        )
        if not splits:
            print(f"[join_layers] No layer-split .gguf files found in {chunk_dir}")
            sys.exit(1)
        paths = [os.path.join(chunk_dir, f) for f in splits]
        join_gguf_layers(paths, args.file)
    elif args.cmd == 'join':
        join_gguf(args.file)
    elif args.cmd == 'status':
        status()
    else:
        p.print_help()


if __name__ == '__main__':
    main()