#!/usr/bin/env python3
"""把 PLE 表的 header 就地改成 Strata 认识的 `strata-ple` 格式（不搬 54 GiB 数据）。

What: 只重写文件头 2912 字节：把 53/50 个键的 header 换成最小 header
      （只有 general.architecture = "strata-ple" + 一个张量描述），再用 0 填充到原数据段起点。
Why:  Strata 的 generate.cpp:4374-4378 只有一句判定：
        ple_only = (general.architecture == "strata-ple")
      只有它为真时，--ple-gguf 才不被追加进 native dense shards、也才跳过
      "additional shard must match … split metadata" 与 "duplicate split shard number" 两条检查。
      官方就是用这个格式：tools/ple_fp8_pack.py 产出 `strata.ple.format` 的 strata-ple 文件。
      社区三个 issue（#3280 / #3875 / #5248）都在请求"放宽到普通 GGUF"，至今全部 open 未合入
      —— 所以走普通 GGUF + qwen4exp 的路必然失败，只有 strata-ple 能过。
How:  张量的 offset 是相对数据段起点的，header 长度不变（用 padding 补回原起点），所以数据段原样不动。
      注意：本脚本只改 header，**数据段仍停在旧位置**（例如 2912）。引擎按最小 header 算出的
      data_start 是 160，两者不一致会报 "PLE table size mismatch ... from data_start 160"，
      所以改完 header 必须再跑一次 ple_move_data.py <file> <旧起点> 160。
用法: python3 ple_to_strata.py <file.gguf>            # 就地改
      python3 ple_to_strata.py <file.gguf> --dry-run  # 只看会写什么
"""
import struct, sys

GGUF_TYPE_STRING = 8
ALIGN = 32


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


def parse(path):
    buf = open(path, "rb").read(1 << 20)
    assert buf[:4] == b"GGUF"
    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
    SZ = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
    for _ in range(n_kv):
        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]
    tens_start = pos
    name_len = struct.unpack_from("<Q", buf, pos)[0]
    pos += 8 + name_len
    nd = struct.unpack_from("<I", buf, pos)[0]; pos += 4
    dims = [struct.unpack_from("<Q", buf, pos + 8 * i)[0] for i in range(nd)]
    pos += 8 * nd
    ttype = struct.unpack_from("<I", buf, pos)[0]; pos += 4
    toff = struct.unpack_from("<Q", buf, pos)[0]; pos += 8
    tens_raw = buf[tens_start:pos]
    tname = buf[tens_start + 8: tens_start + 8 + name_len].decode()
    data_start = (pos + ALIGN - 1) // ALIGN * ALIGN
    return ver, n_t, tname, dims, ttype, toff, tens_raw, data_start, pos


def main():
    path = sys.argv[1]
    dry = "--dry-run" in sys.argv
    import os
    total = os.path.getsize(path)
    ver, n_t, tname, dims, ttype, toff, tens_raw, data_start, hdr_end = parse(path)
    if n_t != 1 or tname != "per_layer_token_embd.weight":
        sys.exit(f"不是单张量 PLE 表（tensors={n_t}, name={tname}）")
    print(f"{path}: tensors=1 {tname} dims={dims} type={ttype} offset={toff} 数据段起点={data_start} 文件={total}")

    header = (b"GGUF" + struct.pack("<I", ver) + struct.pack("<Q", 1) + struct.pack("<Q", 1)
              + wstr("general.architecture") + struct.pack("<I", GGUF_TYPE_STRING) + wstr("strata-ple")
              + tens_raw)
    if len(header) > data_start:
        sys.exit(f"新 header {len(header)} 超过原数据段起点 {data_start}，需重写文件")
    block = header + b"\0" * (data_start - len(header))
    assert len(block) == data_start
    print(f"新 header {len(header)} 字节 + 填充 → 写前 {data_start} 字节（数据段不动）")
    if dry:
        return
    with open(path, "r+b") as f:
        f.write(block)
        f.flush(); os.fsync(f.fileno())
    print("已就地改写 header")


if __name__ == "__main__":
    main()