#!/usr/bin/env python3
"""给 unsloth 的 engram（PLE 表）补齐 Strata 需要的全部架构元数据。

What: 从模型的第一个 GGUF 分片里取出所有 general.* 与 qwen4exp.* 键（原始字节），
      连同 PLE 表自己的 split.* 键，重建成一个新的 PLE GGUF 文件。
Why:  Strata 的 native_dense.cpp:302 对 `--ple-gguf` 的文件调 check_engine_architecture，
      后者经 read_geometry 要求至少六个必备键：
        qwen4exp.block_count / embedding_length / expert_count / expert_used_count /
        attention.head_count / attention.head_count_kv
      缺任何一个就报 "missing qwen4exp.<key>"（实测：只加 general.architecture 后报 block_count）。

      ⚠️ 重要（实测结论，1106）：这条路走不通。把 arch 设成 qwen4exp 后，文件仍会被当成模型的
      第 5 个分片，然后撞上 native_dense.cpp 的 split 检查：模型 4 片已占满 split.no=0..3，
      任何值都会 duplicate（改 0 撞第一片、改 3 撞第四片、不带 split.* 又要求其余键齐全）。
      最终可行方案是 ple_to_strata.py（arch=strata-ple + 无 split.* + 数据段前移到 160）。
      本脚本保留下来只作为"我们试过什么"的记录与字节级 GGUF 头重建的参考实现。
How:  纯字节级重建：header = 原 PLE 的魔数/版本/张量数 + 合并后的 kv + 原张量信息；
      32 字节对齐后原样拷贝张量数据（张量 offset 相对数据段，不受 header 变长影响）。
用法: python3 fix_ple_header2.py <model_shard1.gguf> <ple_src.gguf> <out.gguf>
"""
import os, struct, sys

ALIGN = 32
SZ = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
KEEP_PREFIXES = ("general.", "qwen4exp.")


def align(n, a=ALIGN):
    return (n + a - 1) // a * a


def parse(path):
    """返回 (buf, version, n_tensors, kv_list[(key, raw)], tensor_info_raw, kv_end, tens_end)。"""
    buf = open(path, "rb").read(512 << 20)     # 模型分片有 1224 个张量，header 可达数十 MB
    assert buf[:4] == b"GGUF", path
    pos = 4
    ver = struct.unpack_from("<I", buf, pos)[0]; pos += 4
    n_t = struct.unpack_from("<Q", buf, pos)[0]; pos += 8
    n_kv = struct.unpack_from("<Q", buf, pos)[0]; pos += 8
    kvs = []
    for _ in range(n_kv):
        start = pos
        klen = struct.unpack_from("<Q", buf, pos)[0]; pos += 8
        key = buf[pos:pos + klen].decode(); pos += klen
        vt = struct.unpack_from("<I", buf, pos)[0]; pos += 4
        if vt == 8:
            sl = struct.unpack_from("<Q", buf, pos)[0]; pos += 8 + sl
        elif vt == 9:
            it = struct.unpack_from("<I", buf, pos)[0]; pos += 4
            n = struct.unpack_from("<Q", buf, pos)[0]; pos += 8
            if it == 8:
                for _ in range(n):
                    sl = struct.unpack_from("<Q", buf, pos)[0]; pos += 8 + sl
            else:
                pos += SZ[it] * n
        else:
            pos += SZ[vt]
        kvs.append((key, buf[start:pos]))
    kv_end = pos
    tens_start = pos
    for _ in range(n_t):
        nlen = struct.unpack_from("<Q", buf, pos)[0]; pos += 8 + nlen
        nd = struct.unpack_from("<I", buf, pos)[0]; pos += 4
        pos += 8 * nd + 4 + 8
    tens_end = pos
    return buf, ver, n_t, kvs, buf[tens_start:tens_end], kv_end, tens_end


def main():
    model, ple_src, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
    total = os.path.getsize(ple_src)
    mbuf, mver, mnt, mkvs, _, _, mtend = parse(model)
    pbuf, pver, pnt, pkvs, ptens, pkv_end, ptend = parse(ple_src)

    model_keys = [(k, raw) for k, raw in mkvs if k.startswith(KEEP_PREFIXES)]
    ple_keys = []          # 不复制任何 split.* 键：带了 split.no 必然与模型分片撞号（duplicate）
    # general.architecture 以模型那份为准（qwen4exp）
    merged = dict()
    for k, raw in model_keys:
        merged[k] = raw
    for k, raw in ple_keys:
        merged[k] = raw
    order = list(merged.keys())
    present = [k for k in order if k in ("general.architecture", "general.name")]
    rest = [k for k in order if k not in present]
    body = b"".join(merged[k] for k in present + rest)

    header = (b"GGUF" + struct.pack("<I", pver) + struct.pack("<Q", pnt)
              + struct.pack("<Q", len(merged)) + body + ptens)
    pad = b"\0" * (align(len(header)) - len(header))
    data_old = align(ptend)
    count = total - data_old
    print(f"模型键 {len(model_keys)} 个（不含 split.*，避免与模型分片撞号）= {len(merged)} 个")
    print(f"  header {ptend} → {len(header)} (+pad {len(pad)})；数据 {data_old} → {len(header)+len(pad)}")
    print(f"  拷贝 {count/2**30:.2f} GiB")

    with open(out_path, "wb") as out:
        out.write(header); out.write(pad)
        with open(ple_src, "rb") as fin:
            fin.seek(data_old)
            left = count
            while left > 0:
                chunk = fin.read(min(left, 64 << 20))
                if not chunk:
                    sys.exit("源文件提前结束")
                out.write(chunk); left -= len(chunk)
        out.flush(); os.fsync(out.fileno())
    print(f"完成: {out_path} ({os.path.getsize(out_path)} 字节)")


if __name__ == "__main__":
    main()