#!/usr/bin/env python3
"""Add tokenizer.ggml.mask_token_id to the DSpark draft GGUF.

llama.cpp's DSpark spec reads the draft's mask token via
llama_vocab_mask(vocab) -> tokenizer.ggml.mask_token_id. Our rewritten draft
copied K3's tokenizer, which has no mask token, so it returns -1 and the draft
batch gets fed token -1 -> 'invalid token[1] = -1' -> draft decode fails every step.

The draft's *trained* mask id lives in dflash.mask_token_id (163824). We copy it
into tokenizer.ggml.mask_token_id so llama_vocab_mask() returns the right id.
"""
import struct, sys

SRC = "/root/models/k3-draft/draft_final.gguf"
DST = "/root/models/k3-draft/draft_masked.gguf"
MASK_KEY = b"tokenizer.ggml.mask_token_id"
DFLASH_MASK_KEY = b"dflash.mask_token_id"

def read_kv(data, pos):
    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:
        vlen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
        val = data[pos:pos+vlen]; pos += vlen
    elif vtype in (0, 1, 7):
        val = data[pos:pos+1]; pos += 1
    elif vtype in (2, 3):
        val = data[pos:pos+2]; pos += 2
    elif vtype in (4, 5, 6):
        val = data[pos:pos+4]; pos += 4
    elif vtype in (10, 11, 12):
        val = data[pos:pos+8]; pos += 8
    elif vtype == 9:
        atype = struct.unpack_from("<I", data, pos)[0]; pos += 4
        alen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
        if atype == 8:
            vals = []
            for _ in range(alen):
                slen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
                vals.append(data[pos:pos+slen]); pos += slen
            val = (atype, vals)
        elif atype in (0, 1, 7):
            val = (atype, data[pos:pos+alen]); pos += alen
        elif atype in (2, 3):
            val = (atype, data[pos:pos+alen*2]); pos += alen*2
        elif atype in (4, 5, 6):
            val = (atype, data[pos:pos+alen*4]); pos += alen*4
        elif atype in (10, 11, 12):
            val = (atype, data[pos:pos+alen*8]); pos += alen*8
        else:
            raise ValueError(f"array atype {atype}")
    else:
        raise ValueError(f"vtype {vtype} key {key}")
    return key, vtype, val, pos

def write_kv(out, key, vtype, val):
    out += struct.pack("<Q", len(key)) + key + struct.pack("<I", vtype)
    if vtype == 8:
        out += struct.pack("<Q", len(val)) + val
    elif vtype in (0, 1, 7, 2, 3, 4, 5, 6, 10, 11, 12):
        out += val
    elif vtype == 9:
        atype, d = val
        out += struct.pack("<I", atype)
        if atype == 8:
            out += struct.pack("<Q", len(d))
            for s in d:
                out += struct.pack("<Q", len(s)) + s
        else:
            sz = {0:1,1:1,7:1,2:2,3:2,4:4,5:4,6:4,10:8,11:8,12:8}[atype]
            out += struct.pack("<Q", len(d)//sz) + d

def main():
    data = open(SRC, "rb").read()
    pos = 0
    magic = data[0: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

    kvs = []
    mask_val = None
    have_tok_mask = False
    for _ in range(n_kv):
        key, vtype, val, pos = read_kv(data, pos)
        if key == DFLASH_MASK_KEY:
            mask_val = struct.unpack("<I", val)[0] if vtype in (4,5) else None
            print(f"found {DFLASH_MASK_KEY.decode()} vtype={vtype} val={mask_val}")
        if key == MASK_KEY:
            have_tok_mask = True
            print(f"WARNING: {MASK_KEY.decode()} already present")
        kvs.append((key, vtype, val))

    if have_tok_mask:
        print("mask token already present, nothing to do"); return
    if mask_val is None:
        print("ERROR: dflash.mask_token_id not found"); sys.exit(1)

    # Append the new key
    kvs.append((MASK_KEY, 4, struct.pack("<I", mask_val)))  # uint32
    print(f"adding {MASK_KEY.decode()} = {mask_val} (uint32)")

    # Rebuild header + KV
    out = bytearray()
    out += magic + struct.pack("<I", version) + struct.pack("<Q", n_tensors) + struct.pack("<Q", len(kvs))
    for key, vtype, val in kvs:
        write_kv(out, key, vtype, val)

    kv_end_new = len(out)
    # tensor infos + data follow unchanged from original pos
    rest = data[pos:]

    # The tensor data section is aligned; tensor info offsets are relative to data start.
    # We must recompute alignment: original data start vs new data start.
    alignment = 32
    # original tensor-data start
    orig_data_start = (pos + 0)  # pos == end of KV in original
    # tensor infos occupy from pos; find their end by walking n_tensors
    ti_pos = pos
    for _ in range(n_tensors):
        nl = struct.unpack_from("<Q", data, ti_pos)[0]; ti_pos += 8 + nl
        nd = struct.unpack_from("<I", data, ti_pos)[0]; ti_pos += 4 + nd*8
        ti_pos += 4  # dtype
        ti_pos += 8  # offset
    orig_tensor_data_start = (ti_pos + alignment - 1)//alignment*alignment

    # new tensor-data start
    new_kv_and_ti_len = len(out) + (ti_pos - pos)
    new_tensor_data_start = (new_kv_and_ti_len + alignment - 1)//alignment*alignment

    # append tensor infos verbatim
    out += data[pos:ti_pos]
    while len(out) < new_tensor_data_start:
        out += b"\x00"
    out += data[orig_tensor_data_start:]

    open(DST, "wb").write(out)
    print(f"wrote {DST}: {len(out)} bytes (src {len(data)})")
    print("verify: tokenizer.ggml.mask_token_id should now read", mask_val)

if __name__ == "__main__":
    main()