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


GGML_FILE_MAGIC = 0x67676D6C
GGUF_TYPE_STRING = 8
GGML_TYPE_F32 = 0
GGUF_ALIGNMENT = 32


def pad32(n: int) -> int:
    return (n + GGUF_ALIGNMENT - 1) & ~(GGUF_ALIGNMENT - 1)


def write_gguf_string(buf: bytearray, value: str) -> None:
    data = value.encode("utf-8")
    buf += struct.pack("<Q", len(data))
    buf += data


def tensor_nbytes(dims: tuple[int, ...], ggml_type: int) -> int:
    if ggml_type != GGML_TYPE_F32:
        raise ValueError("only F32 is implemented for these probes")
    n = 1
    for dim in dims:
        n *= dim
    return n * 4


def write_minimal_gguf(path: Path, tensors: list[tuple[str, tuple[int, ...], bytes]], kv: dict[str, str] | None = None) -> None:
    kv = kv or {}
    buf = bytearray()
    buf += b"GGUF"
    buf += struct.pack("<Iqq", 3, len(tensors), len(kv))

    for key, value in kv.items():
        write_gguf_string(buf, key)
        buf += struct.pack("<i", GGUF_TYPE_STRING)
        write_gguf_string(buf, value)

    offset = 0
    for name, dims, data in tensors:
        expected = tensor_nbytes(dims, GGML_TYPE_F32)
        if len(data) != expected:
            raise ValueError(f"{name}: data size {len(data)} != expected {expected}")
        write_gguf_string(buf, name)
        buf += struct.pack("<I", len(dims))
        for dim in dims:
            buf += struct.pack("<q", dim)
        buf += struct.pack("<iQ", GGML_TYPE_F32, offset)
        offset += pad32(expected)

    while len(buf) % GGUF_ALIGNMENT:
        buf += b"\x00"

    for _name, dims, data in tensors:
        buf += data
        while len(buf) % GGUF_ALIGNMENT:
            buf += b"\x00"

    path.write_bytes(buf)


def make_gptj_quantizer_ndims(path: Path, n_dims: int) -> None:
    buf = bytearray()
    buf += struct.pack("<I", GGML_FILE_MAGIC)

    # GPT-J hparams: n_vocab, n_ctx, n_embd, n_head, n_layer, n_rot, ftype.
    buf += struct.pack("<7i", 1, 1, 1, 1, 1, 1, 0)

    # Minimal one-token vocabulary accepted before common tensor parsing.
    token = b"X"
    buf += struct.pack("<iI", 1, len(token))
    buf += token

    name = b"continued3.weight"
    buf += struct.pack("<iii", n_dims, len(name), 0)
    for _ in range(n_dims):
        buf += struct.pack("<i", 1)
    buf += name
    buf += struct.pack("<f", 0.0)
    path.write_bytes(buf)


def make_magika_bad_shapes(path: Path) -> None:
    names = [
        "dense/kernel:0",
        "dense/bias:0",
        "layer_normalization/gamma:0",
        "layer_normalization/beta:0",
        "dense_1/kernel:0",
        "dense_1/bias:0",
        "dense_2/kernel:0",
        "dense_2/bias:0",
        "layer_normalization_1/gamma:0",
        "layer_normalization_1/beta:0",
        "target_label/kernel:0",
        "target_label/bias:0",
    ]
    tensors = [(name, (1,), struct.pack("<f", 0.0)) for name in names]
    write_minimal_gguf(path, tensors)


def make_yolo_empty(path: Path) -> None:
    write_minimal_gguf(path, [])


def make_mnist_bad_arch(path: Path) -> None:
    write_minimal_gguf(path, [], {"general.architecture": "mnist-fc"})


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--outdir", type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)

    make_gptj_quantizer_ndims(args.outdir / "ggml_gptj_quant_common_ndims1024.bin", 1024)
    make_magika_bad_shapes(args.outdir / "magika_required_tensors_bad_shapes.gguf")
    make_yolo_empty(args.outdir / "yolo_empty_model.gguf")
    make_mnist_bad_arch(args.outdir / "mnist_arch_no_tensors.gguf")


if __name__ == "__main__":
    main()