#!/usr/bin/env python3
"""Fix the DSpark draft GGUF:
1. Rename dflash-draft.* keys to dflash.*
2. Rename nested dflash.dflash.* keys to dflash.* equivalents
3. Add tokenizer keys from K3 model
"""
import struct, sys, os

def read_gguf_kv(data, pos, n_kv):
    """Read all KV pairs, return list of (key_bytes, vtype, value_bytes_or_data)"""
    kvs = []
    for i in range(n_kv):
        klen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
        key = data[pos:pos+klen]; pos += klen
        vtype = struct.unpack_from("<I", data, pos)[0]; pos += 4

        if vtype == 8:  # string
            vlen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
            val = data[pos:pos+vlen]; pos += vlen
            kvs.append((key, vtype, val))
        elif vtype in (0, 1):  # u8, i8
            val = data[pos:pos+1]; pos += 1
            kvs.append((key, vtype, val))
        elif vtype in (2, 3):  # u16, i16
            val = data[pos:pos+2]; pos += 2
            kvs.append((key, vtype, val))
        elif vtype in (4, 5, 6):  # u32, i32, f32
            val = data[pos:pos+4]; pos += 4
            kvs.append((key, vtype, val))
        elif vtype == 7:  # bool
            val = data[pos:pos+1]; pos += 1
            kvs.append((key, vtype, val))
        elif vtype in (10, 11, 12):  # u64, i64, f64
            val = data[pos:pos+8]; pos += 8
            kvs.append((key, vtype, val))
        elif vtype == 9:  # array
            atype = struct.unpack_from("<I", data, pos)[0]; pos += 4
            alen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
            if atype == 8:  # string array
                vals = []
                for j in range(alen):
                    slen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
                    sval = data[pos:pos+slen]; pos += slen
                    vals.append(sval)
                kvs.append((key, vtype, (atype, vals)))
            elif atype in (0, 1, 7):  # u8, i8, bool
                val = data[pos:pos+alen]; pos += alen
                kvs.append((key, vtype, (atype, val)))
            elif atype in (2, 3):  # u16, i16
                nbytes = alen * 2
                val = data[pos:pos+nbytes]; pos += nbytes
                kvs.append((key, vtype, (atype, val)))
            elif atype in (4, 5, 6):  # u32, i32, f32
                nbytes = alen * 4
                val = data[pos:pos+nbytes]; pos += nbytes
                kvs.append((key, vtype, (atype, val)))
            elif atype in (10, 11, 12):  # u64, i64, f64
                nbytes = alen * 8
                val = data[pos:pos+nbytes]; pos += nbytes
                kvs.append((key, vtype, (atype, val)))
            else:
                raise ValueError(f"Unknown array type {atype}")
        else:
            raise ValueError(f"Unknown value type {vtype} for key {key}")

    return kvs, pos

def write_gguf_kv(out, key, vtype, val):
    """Write one KV pair to output bytearray"""
    out += struct.pack("<Q", len(key))
    out += key
    out += struct.pack("<I", vtype)

    if vtype == 8:  # string
        out += struct.pack("<Q", len(val))
        out += val
    elif vtype in (0, 1, 7):  # u8, i8, bool
        out += val
    elif vtype in (2, 3):  # u16, i16
        out += val
    elif vtype in (4, 5, 6):  # u32, i32, f32
        out += val
    elif vtype in (10, 11, 12):  # u64, i64, f64
        out += val
    elif vtype == 9:
        atype, data = val
        if atype == 8:  # string array
            out += struct.pack("<I", atype)
            out += struct.pack("<Q", len(data))
            for s in data:
                out += struct.pack("<Q", len(s))
                out += s
        else:
            out += struct.pack("<I", atype)
            if atype in (0, 1, 7):
                alen = len(data)
            elif atype in (2, 3):
                alen = len(data) // 2
            elif atype in (4, 5, 6):
                alen = len(data) // 4
            elif atype in (10, 11, 12):
                alen = len(data) // 8
            out += struct.pack("<Q", alen)
            out += data

def main():
    draft_src = "/root/models/k3-draft/original.gguf"
    k3_src = "/root/models/Kimi-K3-GGUF/UD-Q4_K_XL/Kimi-K3-UD-Q4_K_XL-00001-of-00032.gguf"
    dst = "/root/models/k3-draft/draft_final.gguf"

    # Key renames for dflash-specific keys
    KEY_RENAMES = {
        b"dflash-draft.dflash.target_layer_ids": b"dflash.target_layers",
        b"dflash-draft.dflash.block_size": b"dflash.block_size",
        b"dflash-draft.dflash.n_target_layers": b"dflash.n_target_layers",
        b"dflash-draft.dflash.n_target_features": b"dflash.n_target_features",
        b"dflash-draft.dflash.target.block_count": b"dflash.target_block_count",
        b"dflash-draft.dflash.mask_token_id": b"dflash.mask_token_id",
        b"dflash-draft.dflash.target.repository": b"dflash.target_repository",
        b"dflash-draft.dflash.dspark.enabled": b"dflash.dspark_enabled",
        b"dflash-draft.dflash.dspark.markov_rank": b"dflash.markov_rank",
        b"dflash-draft.dflash.dspark.vocab_size": b"dflash.dspark_vocab_size",
        b"dflash-draft.dflash.dspark.markov_type": b"dflash.markov_type",
        b"dflash-draft.dflash.dspark.confidence.enabled": b"dflash.confidence_enabled",
        b"dflash-draft.dflash.dspark.confidence.with_markov": b"dflash.confidence_with_markov",
        b"dflash-draft.dflash.dspark.confidence_dim": b"dflash.confidence_dim",
    }
    GENERAL_PREFIX_OLD = b"dflash-draft."
    GENERAL_PREFIX_NEW = b"dflash."

    # Read K3 tokenizer keys (stream - file too large to read into memory)
    print("Reading K3 tokenizer keys...")
    with open(k3_src, "rb") as f:
        magic = f.read(4)
        version = struct.unpack("<I", f.read(4))[0]
        n_tensors_k3 = struct.unpack("<Q", f.read(8))[0]
        n_kv_k3 = struct.unpack("<Q", f.read(8))[0]

        k3_kvs = []
        for i in range(n_kv_k3):
            klen = struct.unpack("<Q", f.read(8))[0]
            key = f.read(klen)
            vtype = struct.unpack("<I", f.read(4))[0]

            if vtype == 8:
                vlen = struct.unpack("<Q", f.read(8))[0]
                val = f.read(vlen)
                k3_kvs.append((key, vtype, val))
            elif vtype in (0, 1):  # u8, i8
                val = f.read(1)
                k3_kvs.append((key, vtype, val))
            elif vtype in (2, 3):  # u16, i16
                val = f.read(2)
                k3_kvs.append((key, vtype, val))
            elif vtype in (4, 5, 6):  # u32, i32, f32
                val = f.read(4)
                k3_kvs.append((key, vtype, val))
            elif vtype == 7:  # bool
                val = f.read(1)
                k3_kvs.append((key, vtype, val))
            elif vtype in (10, 11, 12):  # u64, i64, f64
                val = f.read(8)
                k3_kvs.append((key, vtype, val))
            elif vtype == 9:
                atype = struct.unpack("<I", f.read(4))[0]
                alen = struct.unpack("<Q", f.read(8))[0]
                if atype == 8:
                    vals = []
                    for j in range(alen):
                        slen = struct.unpack("<Q", f.read(8))[0]
                        sval = f.read(slen)
                        vals.append(sval)
                    k3_kvs.append((key, vtype, (atype, vals)))
                elif atype in (0, 1, 7):  # u8, i8, bool
                    val = f.read(alen)
                    k3_kvs.append((key, vtype, (atype, val)))
                elif atype in (2, 3):  # u16, i16
                    val = f.read(alen * 2)
                    k3_kvs.append((key, vtype, (atype, val)))
                elif atype in (4, 5, 6):  # u32, i32, f32
                    val = f.read(alen * 4)
                    k3_kvs.append((key, vtype, (atype, val)))
                elif atype in (10, 11, 12):  # u64, i64, f64
                    val = f.read(alen * 8)
                    k3_kvs.append((key, vtype, (atype, val)))

    tok_keys_needed = [
        b"tokenizer.ggml.model",
        b"tokenizer.ggml.pre",
        b"tokenizer.ggml.tokens",
        b"tokenizer.ggml.token_type",
        b"tokenizer.ggml.merges",
    ]
    tok_kvs = {}
    for key, vtype, val in k3_kvs:
        if key in tok_keys_needed:
            tok_kvs[key] = (vtype, val)
            if vtype == 8:
                print(f"  {key.decode()} = {val[:60]}... ({len(val)} bytes)")
            elif vtype == 9:
                atype, adata = val
                print(f"  {key.decode()} = [array type={atype} len={len(adata)}]")
            else:
                print(f"  {key.decode()} = type {vtype}")

    del k3_kvs  # free after extracting tokenizer keys below

    # Read draft GGUF
    print("\nReading draft GGUF...")
    with open(draft_src, "rb") as f:
        data = f.read()

    pos = 0
    magic = data[pos:pos+4]; pos += 4
    version = struct.unpack_from("<I", data, pos)[0]; pos += 4
    n_tensors = struct.unpack_from("<Q", data, pos)[0]; pos += 8
    n_kv = struct.unpack_from("<Q", data, pos)[0]; pos += 8

    draft_kvs, kv_end = read_gguf_kv(data, pos, n_kv)

    # Build new KV list: renamed draft keys + tokenizer keys
    new_kvs = []
    renamed = 0
    existing_keys = set()

    for key, vtype, val in draft_kvs:
        if key in KEY_RENAMES:
            new_key = KEY_RENAMES[key]
            renamed += 1
            print(f"  rename: {key.decode()} -> {new_key.decode()}")
        elif key.startswith(GENERAL_PREFIX_OLD):
            new_key = GENERAL_PREFIX_NEW + key[len(GENERAL_PREFIX_OLD):]
            renamed += 1
        else:
            new_key = key
        new_kvs.append((new_key, vtype, val))
        existing_keys.add(new_key)

    # Add missing tokenizer keys
    added = 0
    for key in tok_keys_needed:
        if key not in existing_keys and key in tok_kvs:
            vtype, val = tok_kvs[key]
            new_kvs.append((key, vtype, val))
            added += 1
            print(f"  added: {key.decode()}")

    print(f"\nRenamed {renamed} keys, added {added} tokenizer keys")
    print(f"Total KV: {len(new_kvs)} (was {n_kv})")

    # Write output
    out = bytearray()
    out += magic
    out += struct.pack("<I", version)
    out += struct.pack("<Q", n_tensors)
    out += struct.pack("<Q", len(new_kvs))

    for key, vtype, val in new_kvs:
        write_gguf_kv(out, key, vtype, val)

    # Read and rewrite tensor infos (rename tensors)
    TENSOR_RENAMES = {
        b"dflash.fc.weight": b"fc.weight",
        b"dflash.hidden_norm.weight": b"enc.output_norm.weight",
        b"dflash.dspark.confidence.weight": b"conf_proj.weight",
        b"dflash.dspark.confidence.bias": b"conf_proj.bias",
        b"dflash.dspark.markov.w1": b"markov_w1.weight",
        b"dflash.dspark.markov.w2": b"markov_w2.weight",
    }

    tensor_infos_start = kv_end
    pos = kv_end
    tensor_infos_out = bytearray()
    tensors_renamed = 0
    for i in range(n_tensors):
        tname_len = struct.unpack_from("<Q", data, pos)[0]; pos += 8
        tname = data[pos:pos+tname_len]; pos += tname_len
        n_dims = struct.unpack_from("<I", data, pos)[0]; pos += 4
        dims = data[pos:pos+n_dims*8]; pos += n_dims * 8
        dtype = data[pos:pos+4]; pos += 4
        offset = data[pos:pos+8]; pos += 8

        # Rename tensor if needed
        if tname in TENSOR_RENAMES:
            new_tname = TENSOR_RENAMES[tname]
            tensors_renamed += 1
            print(f"  tensor rename: {tname.decode()} -> {new_tname.decode()}")
        else:
            new_tname = tname

        tensor_infos_out += struct.pack("<Q", len(new_tname))
        tensor_infos_out += new_tname
        tensor_infos_out += struct.pack("<I", n_dims)
        tensor_infos_out += dims
        tensor_infos_out += dtype
        tensor_infos_out += offset

    tensor_infos_end = pos
    print(f"Renamed {tensors_renamed} tensors")

    alignment = 32
    new_tensor_data_start = (len(out) + len(tensor_infos_out) + alignment - 1) // alignment * alignment
    orig_tensor_data_start = (tensor_infos_end + alignment - 1) // alignment * alignment

    out += tensor_infos_out
    while len(out) < new_tensor_data_start:
        out += b"\x00"
    out += data[orig_tensor_data_start:]

    with open(dst, "wb") as f:
        f.write(out)

    print(f"\nWrote {dst}: {len(out)} bytes (orig: {len(data)})")

if __name__ == "__main__":
    main()