#!/usr/bin/env python
"""Splice replacement tensors into a split GGUF, writing a new shard set. Pure Python, no numpy.

Every tensor of the original shards is copied byte-for-byte unless a tensor of the same name exists
in the replacement GGUF, in which case the replacement's shape/type/data is used instead. Shard 1
(metadata only in gguf-split layouts) is copied verbatim under the new name; shards with tensors get
their split.* keys regenerated and tensor offsets recomputed. The originals are never modified.

Usage:
  python gguf_splice.py <original first shard> <replacement.gguf> <output base name>
  -> writes <output base name>-0000N-of-0000M.gguf next to the original shards
Example (2026-09-03):
  python gguf_splice.py Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf build-downs-iq4nl.gguf Qwen3.8-Flash-Next-UD-Q4_K_XL-dn4

Verify afterwards with gguf_tensor_map.py on the new first shard (tensor count must match split.tensors.count).
"""
import os, re, shutil, struct, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gguf_extract import Reader, nbytes, wstr, shards_of, ALIGN


def write_shard(out_path, kv, tensors, sources):
    """kv: list of (key, type, raw); tensors: list of (name, shape, type); sources: name -> (Reader, src_off)."""
    hdr = b'GGUF' + struct.pack('<IQQ', 3, len(tensors), len(kv))
    for key, ty, raw in kv: hdr += wstr(key) + struct.pack('<I', ty) + raw
    infos = b''; off = 0; layout = []
    for name, shape, ty in tensors:
        sz = nbytes(shape, ty); layout.append((name, shape, ty, off, sz))
        infos += wstr(name) + struct.pack('<I', len(shape)) + b''.join(struct.pack('<Q', d) for d in shape) + struct.pack('<IQ', ty, off)
        off += (sz + ALIGN - 1) // ALIGN * ALIGN
    head = hdr + infos
    pad = (ALIGN - len(head) % ALIGN) % ALIGN
    with open(out_path, 'wb') as o:
        o.write(head + b'\x00' * pad)
        for name, shape, ty, dst_off, sz in layout:
            r, src_off = sources[name]
            r.f.seek(r.data_start + src_off); remaining = sz
            while remaining:
                chunk = r.f.read(min(remaining, 64 << 20)); o.write(chunk); remaining -= len(chunk)
            o.write(b'\x00' * ((ALIGN - sz % ALIGN) % ALIGN))
    return off


def main():
    first, repl_path, out_base = sys.argv[1], sys.argv[2], sys.argv[3]
    shard_paths = shards_of(first); n = len(shard_paths)
    repl = Reader(repl_path)
    repl_map = {t[0]: t for t in repl.tensors}
    print(f'replacement carries {len(repl_map)} tensors')
    total = 0; replaced = 0; out_dir = os.path.dirname(first) or '.'
    for i, p in enumerate(shard_paths, 1):
        out_path = os.path.join(out_dir, f'{out_base}-{i:05d}-of-{n:05d}.gguf')
        r = Reader(p)
        if r.n_tensors == 0:
            r.f.close(); shutil.copyfile(p, out_path); print(f'shard {i}: metadata only, copied'); continue
        kv = []
        for key, ty, raw in r.kv:
            kv.append((key, ty, raw))   # split.no/count/tensors.count unchanged: same shard numbering, same total
        tensors = []; sources = {}
        for name, shape, ty, off in r.tensors:
            if name in repl_map:
                _, rshape, rty, roff = repl_map[name]
                assert rshape == shape, (name, shape, rshape)
                tensors.append((name, rshape, rty)); sources[name] = (repl, roff); replaced += 1
            else:
                tensors.append((name, shape, ty)); sources[name] = (r, off)
        written = write_shard(out_path, kv, tensors, sources); total += written
        print(f'shard {i}: {len(tensors)} tensors, {written / 1e9:.2f} GB -> {os.path.basename(out_path)}')
        r.f.close()
    print(f'done: {replaced} tensors replaced, {total / 1e9:.2f} GB of tensor data written')
    if replaced != len(repl_map): print(f'WARNING: {len(repl_map) - replaced} replacement tensors were not found in the original')


if __name__ == '__main__':
    main()