#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""surgery3.py — 手术 III: qwen35.attention.recurrent_layers BOOL 数组 33→32 (Colab 版)
沿袭 R10 s8_f16_surgery2.py 实战逻辑: 删末元素(幻影 nextn 层旗标)+计数 33→32+头部左移 1,
align32 padding 吸收位移 → 张量数据区起点不变。
前置: r12_surgery.py 已完成 (bc=32, nx=0)
产物: SURGERY3_DONE"""
import struct, os, sys

PATH = sys.argv[1] if len(sys.argv) > 1 else "/content/k8b/base_f16.gguf"
KEY = "qwen35.attention.recurrent_layers"
SZ = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}

def walk_header(fh):
    fh.seek(0)
    assert fh.read(4) == b"GGUF"
    ver, n_ten, n_kv = struct.unpack("<IQQ", fh.read(20))
    def rstr():
        n = struct.unpack("<Q", fh.read(8))[0]
        return fh.read(n).decode("utf-8", "replace")
    hits = {}
    for _ in range(n_kv):
        key = rstr()
        vtype = struct.unpack("<I", fh.read(4))[0]
        if vtype == 8:
            rstr(); continue
        if vtype == 9:
            etype = struct.unpack("<I", fh.read(4))[0]
            cnt_off = fh.tell()
            cnt = struct.unpack("<Q", fh.read(8))[0]
            data_off = fh.tell()
            if etype == 8:
                for _ in range(cnt):
                    ln = struct.unpack("<Q", fh.read(8))[0]
                    fh.seek(ln, 1)
            else:
                fh.seek(cnt * SZ[etype], 1)
            if key == KEY:
                hits[key] = (vtype, etype, cnt, cnt_off, data_off)
            continue
        fh.seek(SZ[vtype], 1)
    for _ in range(n_ten):
        rstr()
        n_dims = struct.unpack("<I", fh.read(4))[0]
        fh.seek(8 * n_dims, 1)
        fh.seek(4 + 8, 1)
    head_end = fh.tell()
    data_start = (head_end + 31) // 32 * 32
    return hits, head_end, data_start, n_ten

def main():
    sz = os.path.getsize(PATH)
    with open(PATH, "r+b") as fh:
        hits, head_end, data_start, n_ten = walk_header(fh)
        if KEY not in hits:
            print("SKIP: %s absent (maybe already 32 or not present)" % KEY)
            print("SURGERY3_DONE")
            return
        vtype, etype, cnt, cnt_off, data_off = hits[KEY]
        print("found: etype=%d count=%d cnt_off=%d data_off=%d head_end=%d data_start=%d filesize=%d"
              % (etype, cnt, cnt_off, data_off, head_end, data_start, sz), flush=True)
        if cnt == 32:
            print("SKIP: already 32")
            print("SURGERY3_DONE")
            return
        assert etype == 7 and cnt == 33, "expect BOOL array of 33, got %d/%d" % (etype, cnt)
        fh.seek(0)
        hdr = bytearray(fh.read(data_start))
        struct.pack_into("<Q", hdr, cnt_off, 32)
        del hdr[data_off + 32]
        assert len(hdr) == data_start - 1
        hdr.append(0)
        assert len(hdr) == data_start
        fh.seek(0)
        fh.write(hdr)
        fh.flush()
        os.fsync(fh.fileno())
    print("bytes rewritten:", data_start, flush=True)
    with open(PATH, "rb") as fh:
        hits2, head_end2, data_start2, n_ten2 = walk_header(fh)
        print("VERIFY recount=%d head_end=%d data_start=%d (expect unchanged %d)"
              % (hits2[KEY][2], head_end2, data_start2, data_start), flush=True)
        assert hits2[KEY][2] == 32 and data_start2 == data_start
    import gguf
    r = gguf.GGUFReader(PATH)
    print("VERIFY reader: tensors=%d block_count=%s" % (
        len(r.tensors), int(r.fields["qwen35.block_count"].parts[-1][0])), flush=True)
    print("SURGERY3_DONE", flush=True)

if __name__ == "__main__":
    main()