#!/usr/bin/env python3
"""
Streaming GGUF expert slicer for Qwen3.8-Flash-Next (qwen4exp).
Slices blk.{l}.ffn_{down,gate,up}_exps.weight and blk.{l}.ffn_gate_inp.weight
from N_experts (512) down to K experts per layer using a per-layer keep list,
updates qwen4exp.expert_count = K, and streams non-expert tensors via copy_file_range/chunks
with zero re-quantization and no enum dependency on gguf-py.
"""
import argparse
import json
import os
import re
import struct
import sys
import tempfile
from pathlib import Path

GGUF_MAGIC = 0x46554747  # "GGUF" little-endian
DEFAULT_ALIGNMENT = 32

# GGUF metadata value types
GGUF_TYPE_UINT8   = 0
GGUF_TYPE_INT8    = 1
GGUF_TYPE_UINT16  = 2
GGUF_TYPE_INT16   = 3
GGUF_TYPE_UINT32  = 4
GGUF_TYPE_INT32   = 5
GGUF_TYPE_FLOAT32 = 6
GGUF_TYPE_BOOL    = 7
GGUF_TYPE_STRING  = 8
GGUF_TYPE_ARRAY   = 9
GGUF_TYPE_UINT64  = 10
GGUF_TYPE_INT64   = 11
GGUF_TYPE_FLOAT64 = 12

SCALAR_SIZES = {
    GGUF_TYPE_UINT8: 1,
    GGUF_TYPE_INT8: 1,
    GGUF_TYPE_UINT16: 2,
    GGUF_TYPE_INT16: 2,
    GGUF_TYPE_UINT32: 4,
    GGUF_TYPE_INT32: 4,
    GGUF_TYPE_FLOAT32: 4,
    GGUF_TYPE_BOOL: 1,
    GGUF_TYPE_UINT64: 8,
    GGUF_TYPE_INT64: 8,
    GGUF_TYPE_FLOAT64: 8,
}

EXPERT_TENSOR_RE = re.compile(
    r"^blk\.(\d+)\.(ffn_(?:down|gate|up|gate_up)_exps\.(?:weight|bias|scale)|ffn_gate_inp\.(?:weight|bias)|exp_probs_b\.bias)$"
)
NON_EXPERT_512_SUFFIXES = (".attn_k.weight", ".attn_v.weight", ".indexer.q_proj.weight")


def pwrite_all(fd: int, data: bytes | memoryview, offset: int) -> None:
    mv = memoryview(data)
    written = 0
    total = len(mv)
    while written < total:
        n = os.pwrite(fd, mv[written:], offset + written)
        if n <= 0:
            raise OSError(f"pwrite returned {n} at offset {offset + written}")
        written += n


def align_offset(offset: int, alignment: int = DEFAULT_ALIGNMENT) -> int:
    rem = offset % alignment
    return offset if rem == 0 else offset + (alignment - rem)


class BinaryCursor:
    def __init__(self, fd: int):
        self.fd = fd
        self.pos = 0
        self.buf = b""
        self.buf_start = 0

    def _ensure(self, n: int) -> bytes:
        if self.pos + n > self.buf_start + len(self.buf):
            read_len = max(n, 256 * 1024)
            chunk = os.pread(self.fd, read_len, self.pos)
            if len(chunk) < n:
                raise EOFError(f"Unexpected EOF at {self.pos}, wanted {n} bytes")
            self.buf = chunk
            self.buf_start = self.pos
        off = self.pos - self.buf_start
        res = self.buf[off : off + n]
        self.pos += n
        return res

    def read_u32(self) -> int:
        return struct.unpack("<I", self._ensure(4))[0]

    def read_u64(self) -> int:
        return struct.unpack("<Q", self._ensure(8))[0]

    def read_string_raw(self) -> tuple[str, bytes]:
        raw_len = self._ensure(8)
        slen = struct.unpack("<Q", raw_len)[0]
        raw_bytes = self._ensure(slen)
        return raw_bytes.decode("utf-8", errors="replace"), raw_len + raw_bytes

    def read_value_raw(self, vtype: int) -> bytes:
        if vtype in SCALAR_SIZES:
            return self._ensure(SCALAR_SIZES[vtype])
        if vtype == GGUF_TYPE_STRING:
            _, raw = self.read_string_raw()
            return raw
        if vtype == GGUF_TYPE_ARRAY:
            hdr = self._ensure(12)
            elem_type, count = struct.unpack("<IQ", hdr)
            if elem_type in SCALAR_SIZES:
                body = self._ensure(SCALAR_SIZES[elem_type] * count)
                return hdr + body
            parts = [hdr]
            for _ in range(count):
                parts.append(self.read_value_raw(elem_type))
            return b"".join(parts)
        raise ValueError(f"Unknown GGUF metadata value type: {vtype}")


def parse_gguf_header(fd: int, file_size: int) -> dict:
    cur = BinaryCursor(fd)
    magic = cur.read_u32()
    if magic != GGUF_MAGIC:
        raise ValueError(f"Not a GGUF file (magic={hex(magic)})")
    version = cur.read_u32()
    n_tensors = cur.read_u64()
    n_kv = cur.read_u64()

    alignment = DEFAULT_ALIGNMENT
    expert_count = None
    kvs = []
    for _ in range(n_kv):
        key, raw_key = cur.read_string_raw()
        raw_vtype = cur._ensure(4)
        vtype = struct.unpack("<I", raw_vtype)[0]
        raw_val = cur.read_value_raw(vtype)
        if key == "general.alignment" and vtype == GGUF_TYPE_UINT32:
            alignment = struct.unpack("<I", raw_val)[0]
        elif key.endswith(".expert_count") and vtype == GGUF_TYPE_UINT32:
            expert_count = struct.unpack("<I", raw_val)[0]
        kvs.append({
            "key": key,
            "raw_key": raw_key,
            "vtype": vtype,
            "raw_vtype": raw_vtype,
            "raw_val": raw_val,
        })

    tensors = []
    for _ in range(n_tensors):
        name, raw_name = cur.read_string_raw()
        n_dims = cur.read_u32()
        dims = [cur.read_u64() for _ in range(n_dims)]
        ggml_type = cur.read_u32()
        rel_offset = cur.read_u64()
        tensors.append({
            "name": name,
            "raw_name": raw_name,
            "n_dims": n_dims,
            "dims": dims,
            "ggml_type": ggml_type,
            "rel_offset": rel_offset,
        })

    data_section_offset = align_offset(cur.pos, alignment)

    # Determine exact byte size of each tensor from sorted offsets and file size
    # Note: tensors in GGUF are laid out in ascending rel_offset order, each padded to alignment
    for i, t in enumerate(tensors):
        abs_start = data_section_offset + t["rel_offset"]
        t["abs_start"] = abs_start
        if i + 1 < len(tensors):
            next_start = data_section_offset + tensors[i + 1]["rel_offset"]
            # Tensor size is next_start minus alignment padding; for expert tensors we compute exact slice size
            t["padded_span"] = next_start - abs_start
        else:
            t["padded_span"] = file_size - abs_start

    return {
        "version": version,
        "n_tensors": n_tensors,
        "n_kv": n_kv,
        "alignment": alignment,
        "expert_count": expert_count,
        "kvs": kvs,
        "tensors": tensors,
        "data_section_offset": data_section_offset,
    }


# Exact block size (in elements) and byte size per block for all GGML types used in GGUFs
GGML_TYPE_TRAITS = {
    0:  (1, 4),     # F32
    1:  (1, 2),     # F16
    2:  (32, 18),   # Q4_0
    3:  (32, 20),   # Q4_1
    6:  (32, 22),   # Q5_0
    7:  (32, 24),   # Q5_1
    8:  (32, 34),   # Q8_0
    9:  (32, 36),   # Q8_1
    10: (256, 84),  # Q2_K
    11: (256, 110), # Q3_K
    12: (256, 144), # Q4_K
    13: (256, 176), # Q5_K
    14: (256, 210), # Q6_K
    15: (256, 292), # Q8_K
    16: (256, 66),  # IQ2_XXS
    17: (256, 74),  # IQ2_XS
    18: (256, 98),  # IQ3_XXS
    19: (256, 50),  # IQ1_S
    20: (32, 18),   # IQ4_NL
    21: (256, 110), # IQ3_S
    22: (256, 82),  # IQ2_S
    23: (256, 136), # IQ4_XS
    24: (1, 1),     # I8
    25: (1, 2),     # I16
    26: (1, 4),     # I32
    27: (1, 8),     # I64
    28: (1, 8),     # F64
    29: (256, 56),  # IQ1_M
    30: (1, 2),     # BF16
    34: (256, 54),  # TQ1_0
    35: (256, 66),  # TQ2_0
    39: (32, 17),   # MXFP4
    40: (64, 36),   # NVFP4
    41: (128, 18),  # Q1_0
    42: (64, 18),   # Q2_0 (64 weights = 16B qs + 2B fp16 d = 18B, 2.25 bpw)
}


def compute_tensor_nbytes(dims: list[int], ggml_type: int) -> int:
    if ggml_type not in GGML_TYPE_TRAITS:
        raise ValueError(f"Unsupported ggml_type={ggml_type}")
    blck_size, type_size = GGML_TYPE_TRAITS[ggml_type]
    n_elems = 1
    for d in dims:
        n_elems *= d
    if n_elems % blck_size != 0:
        raise ValueError(f"Tensor dims {dims} not divisible by block size {blck_size} for type {ggml_type}")
    return (n_elems // blck_size) * type_size


def copy_range_fast(src_fd: int, src_off: int, dst_fd: int, dst_off: int, length: int) -> None:
    """Fast streaming copy using copy_file_range on Linux or 64 MiB pread/pwrite chunks."""
    has_cfr = hasattr(os, "copy_file_range")
    rem = length
    cur_s = src_off
    cur_d = dst_off
    chunk_sz = 64 * 1024 * 1024
    while rem > 0:
        step = min(rem, chunk_sz)
        if has_cfr:
            try:
                n = os.copy_file_range(src_fd, dst_fd, step, offset_src=cur_s, offset_dst=cur_d)
                if n > 0:
                    cur_s += n
                    cur_d += n
                    rem -= n
                    continue
                has_cfr = False
            except OSError:
                has_cfr = False
        buf = os.pread(src_fd, step, cur_s)
        if not buf:
            raise EOFError(f"Unexpected EOF copying {step} bytes at {cur_s}")
        pwrite_all(dst_fd, buf, cur_d)
        cur_s += len(buf)
        cur_d += len(buf)
        rem -= len(buf)


def slice_gguf_shard(
    in_path: str,
    out_path: str,
    keep_map: dict[int, list[int]],
    old_n_experts: int = 512,
) -> dict:
    if os.path.abspath(in_path) == os.path.abspath(out_path):
        raise ValueError("in_path and out_path must be distinct files")
    file_size = os.path.getsize(in_path)
    in_fd = os.open(in_path, os.O_RDONLY)
    try:
        hdr = parse_gguf_header(in_fd, file_size)
        alignment = hdr["alignment"]

        # Validate keep_map has uniform K across all present layers
        ks = {len(v) for v in keep_map.values()}
        if len(ks) != 1:
            raise ValueError(f"GGUF requires uniform K across layers, got lengths {ks}")
        new_k = next(iter(ks))
        for l, idx_list in keep_map.items():
            if len(set(idx_list)) != new_k or any(i < 0 or i >= old_n_experts for i in idx_list):
                raise ValueError(f"Invalid expert indices at layer {l}")

        # Build updated KV section
        kv_bytes_list = []
        for kv in hdr["kvs"]:
            raw_val = kv["raw_val"]
            if kv["key"].endswith(".expert_used_count") and kv["vtype"] == GGUF_TYPE_UINT32:
                used_k = struct.unpack("<I", raw_val)[0]
                if used_k > new_k:
                    raise ValueError(f"{kv['key']}={used_k} exceeds new expert count K={new_k}")
            if kv["key"].endswith(".expert_count") and kv["vtype"] == GGUF_TYPE_UINT32:
                raw_val = struct.pack("<I", new_k)
            kv_bytes_list.append(kv["raw_key"] + kv["raw_vtype"] + raw_val)

        # Build updated tensor infos and compute new relative offsets
        new_tensors = []
        cur_rel_offset = 0
        sliced_count = 0
        bytes_saved = 0

        for t in hdr["tensors"]:
            dims = list(t["dims"])
            orig_nbytes = compute_tensor_nbytes(dims, t["ggml_type"])
            m = EXPERT_TENSOR_RE.match(t["name"])
            is_expert_tensor = False
            layer_idx = None

            if m is not None and dims[-1] == old_n_experts:
                layer_idx = int(m.group(1))
                if layer_idx not in keep_map:
                    raise KeyError(f"Layer {layer_idx} missing from keep_map for tensor {t['name']}")
                dims[-1] = new_k
                is_expert_tensor = True
                sliced_count += 1
            elif t["name"].startswith("blk.") and dims[-1] == old_n_experts and not t["name"].endswith(NON_EXPERT_512_SUFFIXES):
                raise ValueError(f"Unhandled blk tensor with outermost dim {old_n_experts}: {t['name']} shape={dims}")

            new_nbytes = compute_tensor_nbytes(dims, t["ggml_type"])
            bytes_saved += (orig_nbytes - new_nbytes)

            new_tensors.append({
                "name": t["name"],
                "raw_name": t["raw_name"],
                "n_dims": t["n_dims"],
                "orig_dims": t["dims"],
                "new_dims": dims,
                "ggml_type": t["ggml_type"],
                "orig_abs_start": t["abs_start"],
                "orig_nbytes": orig_nbytes,
                "new_rel_offset": cur_rel_offset,
                "new_nbytes": new_nbytes,
                "is_expert_tensor": is_expert_tensor,
                "layer_idx": layer_idx,
            })
            cur_rel_offset = align_offset(cur_rel_offset + new_nbytes, alignment)

        # Serialize header + KV + tensor directory
        out_hdr_parts = [
            struct.pack("<IIQQ", GGUF_MAGIC, hdr["version"], hdr["n_tensors"], hdr["n_kv"]),
            *kv_bytes_list,
        ]
        for nt in new_tensors:
            dims_packed = struct.pack(f"<{nt['n_dims']}Q", *nt["new_dims"])
            info_packed = (
                nt["raw_name"]
                + struct.pack("<I", nt["n_dims"])
                + dims_packed
                + struct.pack("<IQ", nt["ggml_type"], nt["new_rel_offset"])
            )
            out_hdr_parts.append(info_packed)

        raw_header = b"".join(out_hdr_parts)
        new_data_offset = align_offset(len(raw_header), alignment)
        pad_len = new_data_offset - len(raw_header)
        if pad_len > 0:
            raw_header += b"\x00" * pad_len

        out_fd = os.open(out_path, os.O_CREAT | os.O_TRUNC | os.O_WRONLY, 0o644)
        try:
            pwrite_all(out_fd, raw_header, 0)
            for nt in new_tensors:
                dst_start = new_data_offset + nt["new_rel_offset"]
                if not nt["is_expert_tensor"]:
                    copy_range_fast(in_fd, nt["orig_abs_start"], out_fd, dst_start, nt["orig_nbytes"])
                else:
                    # Gather K kept expert slices via per-expert pread (safe for tensors >2 GiB)
                    per_exp_bytes = nt["orig_nbytes"] // old_n_experts
                    keep_idxs = keep_map[nt["layer_idx"]]
                    sliced_chunks = []
                    for idx in keep_idxs:
                        exp_off = nt["orig_abs_start"] + idx * per_exp_bytes
                        chunk = os.pread(in_fd, per_exp_bytes, exp_off)
                        if len(chunk) != per_exp_bytes:
                            raise EOFError(f"Short read on {nt['name']} expert {idx}")
                        sliced_chunks.append(chunk)
                    out_bytes = b"".join(sliced_chunks)
                    assert len(out_bytes) == nt["new_nbytes"]
                    pwrite_all(out_fd, out_bytes, dst_start)

                # Zero-pad to alignment boundary
                aligned_end = new_data_offset + align_offset(nt["new_rel_offset"] + nt["new_nbytes"], alignment)
                actual_end = dst_start + nt["new_nbytes"]
                if aligned_end > actual_end:
                    pwrite_all(out_fd, b"\x00" * (aligned_end - actual_end), actual_end)
        finally:
            os.close(out_fd)

        return {
            "in_path": in_path,
            "out_path": out_path,
            "old_n_experts": old_n_experts,
            "new_n_experts": new_k,
            "sliced_tensors": sliced_count,
            "bytes_saved": bytes_saved,
            "out_file_size": os.path.getsize(out_path),
        }
    finally:
        os.close(in_fd)


def _pack_str(s: str) -> bytes:
    b = s.encode("utf-8")
    return struct.pack("<Q", len(b)) + b


def run_self_test() -> None:
    """Create a synthetic qwen4exp GGUF with Q2_0 (type 42) 3D expert tensors, slice 8 -> 4 experts, verify."""
    with tempfile.TemporaryDirectory() as td:
        src = os.path.join(td, "test_in.gguf")
        dst = os.path.join(td, "test_out.gguf")

        old_e = 8
        new_e = 4
        keep = {0: [1, 3, 4, 7]}

        # Build synthetic GGUF:
        # KV 1: general.architecture = "qwen4exp"
        # KV 2: qwen4exp.expert_count = 8 (uint32)
        kv1 = _pack_str("general.architecture") + struct.pack("<I", GGUF_TYPE_STRING) + _pack_str("qwen4exp")
        kv2 = _pack_str("qwen4exp.expert_count") + struct.pack("<I", GGUF_TYPE_UINT32) + struct.pack("<I", old_e)

        # Tensor 0: token_embd.weight [64, 16] F32
        t0_dims = [64, 16]
        t0_bytes = os.urandom(compute_tensor_nbytes(t0_dims, 0))
        # Tensor 1: blk.0.ffn_down_exps.weight [64, 32, 8] Q2_0 (type 42: 32 elems = 10B)
        t1_dims = [64, 32, old_e]
        t1_bytes_per_exp = (64 * 32 // 64) * 18
        t1_exp_slices = [os.urandom(t1_bytes_per_exp) for _ in range(old_e)]
        t1_bytes = b"".join(t1_exp_slices)
        # Tensor 2: blk.0.ffn_gate_inp.weight [64, 8] F32
        t2_dims = [64, old_e]
        t2_bytes_per_exp = 64 * 4
        t2_exp_slices = [os.urandom(t2_bytes_per_exp) for _ in range(old_e)]
        t2_bytes = b"".join(t2_exp_slices)

        off0 = 0
        off1 = align_offset(off0 + len(t0_bytes))
        off2 = align_offset(off1 + len(t1_bytes))

        ti0 = _pack_str("token_embd.weight") + struct.pack("<I2QIQ", 2, *t0_dims, 0, off0)
        ti1 = _pack_str("blk.0.ffn_down_exps.weight") + struct.pack("<I3QIQ", 3, *t1_dims, 42, off1)
        ti2 = _pack_str("blk.0.ffn_gate_inp.weight") + struct.pack("<I2QIQ", 2, *t2_dims, 0, off2)

        hdr = struct.pack("<IIQQ", GGUF_MAGIC, 3, 3, 2) + kv1 + kv2 + ti0 + ti1 + ti2
        data_off = align_offset(len(hdr))
        with open(src, "wb") as f:
            f.write(hdr + b"\x00" * (data_off - len(hdr)))
            f.seek(data_off + off0)
            f.write(t0_bytes)
            f.seek(data_off + off1)
            f.write(t1_bytes)
            f.seek(data_off + off2)
            f.write(t2_bytes)

        res = slice_gguf_shard(src, dst, keep, old_n_experts=old_e)
        assert res["sliced_tensors"] == 2, res

        out_fd = os.open(dst, os.O_RDONLY)
        try:
            parsed = parse_gguf_header(out_fd, os.path.getsize(dst))
            assert parsed["expert_count"] == new_e
            doff = parsed["data_section_offset"]
            # Check t0 bit-identical
            got_t0 = os.pread(out_fd, len(t0_bytes), doff + parsed["tensors"][0]["rel_offset"])
            assert got_t0 == t0_bytes
            # Check t1 sliced Q2_0 experts [1, 3, 4, 7] bit-identical
            expected_t1 = b"".join(t1_exp_slices[i] for i in keep[0])
            got_t1 = os.pread(out_fd, len(expected_t1), doff + parsed["tensors"][1]["rel_offset"])
            assert got_t1 == expected_t1
            # Check t2 sliced router rows [1, 3, 4, 7] bit-identical
            expected_t2 = b"".join(t2_exp_slices[i] for i in keep[0])
            got_t2 = os.pread(out_fd, len(expected_t2), doff + parsed["tensors"][2]["rel_offset"])
            assert got_t2 == expected_t2
        finally:
            os.close(out_fd)
    print("slice_qwen4exp_gguf.py self-test: PASS")


def main():
    p = argparse.ArgumentParser(description="Slice experts from Qwen3.8-Flash-Next GGUF shards.")
    p.add_argument("--self-test", action="store_true", help="Run built-in unit test and exit.")
    p.add_argument("--in-gguf", help="Input GGUF shard path.")
    p.add_argument("--out-gguf", help="Output GGUF shard path.")
    p.add_argument("--keep-json", help="JSON file mapping layer index ('0'..'47') to list of kept expert IDs.")
    p.add_argument("--old-experts", type=int, default=512)
    args = p.parse_args()

    if args.self_test:
        run_self_test()
        return

    if not (args.in_gguf and args.out_gguf and args.keep_json):
        p.error("--in-gguf, --out-gguf, and --keep-json are required unless --self-test is set")

    raw_map = json.loads(Path(args.keep_json).read_text())
    keep_map = {int(k): sorted(int(x) for x in v) for k, v in raw_map.items()}
    res = slice_gguf_shard(args.in_gguf, args.out_gguf, keep_map, old_n_experts=args.old_experts)
    print(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()