"""Round-trips every backend test model through the onnxruntime flatbuffer format.

For each backend test case shipped with ``onnx-light`` (enumerated by
:func:`onnx_light.onnx_lib.backend.test.case.collect_test_case`, which is
also the source used by :func:`make_test_class`) the test performs two
round-trips:

1. **onnx-light → onnx-light** -- the model is serialized with
   ``SerializeFormat.ORT_FLATBUFFERS`` by onnx-light, parsed back by
   onnx-light, and the parsed ``ModelProto`` must be byte-identical to the
   original (when re-serialized in the ONNX protobuf format).

2. **onnxruntime → onnx-light** -- the model is loaded by ``onnxruntime``
   and re-saved by ``onnxruntime`` in its native ``.ort`` flatbuffer format
   (via ``SessionOptions.optimized_model_filepath`` with
   ``session.save_model_format=ORT`` and graph optimizations disabled), and
   the resulting file is parsed back by onnx-light. The parsed model must
   again be byte-identical (in protobuf form) to the original.

While ``SerializeFormat.ORT_FLATBUFFERS`` is not implemented in onnx-light
the first step raises ``RuntimeError`` immediately and the whole case is
reported as skipped. This is intentional: the script is the harness that
will pick up the new code path the moment the C++ writer/reader lands,
without any further changes. The second step relies on onnx-light being
able to *read* a ``.ort`` file produced by onnxruntime, so it skips for the
same reason today.

The test methods are generated by :func:`make_test_class`, which produces
one ``test_<case_name>`` per backend test case.
"""

from __future__ import annotations

import os
import tempfile
import unittest

import onnxruntime as ort
from onnx_light.ext_test_case import import_or_skip

import onnx_light.onnx as onnxl

# The backend test registries are only available in the full build; skip this
# module on a reduced build (ONNX_LIGHT_BUILD_KERNELS=OFF).
make_test_class = import_or_skip("onnx_light.onnx.backend", "make_test_class")


def _onnx_bytes(model: onnxl.ModelProto) -> bytes:
    """Returns the canonical ONNX-protobuf serialization of *model*."""
    sopts = onnxl.SerializeOptions()
    sopts.format = onnxl.SerializeFormat.ONNX
    return model.SerializeToString(sopts)


def _save_with_onnx_light_ort(model: onnxl.ModelProto, path: str) -> None:
    """Saves *model* at *path* using onnx-light's ORT flatbuffer writer."""
    sopts = onnxl.SerializeOptions()
    sopts.format = onnxl.SerializeFormat.ORT_FLATBUFFERS
    model.SerializeToFile(path, sopts)


def _load_with_onnx_light_ort(path: str) -> onnxl.ModelProto:
    """Loads an ORT flatbuffer model from *path* using onnx-light."""
    popts = onnxl.ParseOptions()
    popts.format = onnxl.SerializeFormat.ORT_FLATBUFFERS
    parsed = onnxl.ModelProto()
    parsed.ParseFromFile(path, popts)
    return parsed


def _save_with_onnxruntime_ort(src_path: str, dst_path: str) -> None:
    """Re-saves a model as a ``.ort`` flatbuffer file using onnxruntime.

    Loads ``src_path`` (an ``.onnx`` file) into an ``InferenceSession`` with
    graph optimizations disabled so that the serialized graph stays
    structurally equivalent to the input, then asks onnxruntime to dump the
    optimized (== original) model at ``dst_path`` using the ORT flatbuffer
    format.
    """
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    so.optimized_model_filepath = dst_path
    so.add_session_config_entry("session.save_model_format", "ORT")
    # Creating the session triggers the optimized-model dump.
    ort.InferenceSession(src_path, so, providers=["CPUExecutionProvider"])


def ort_round_trip_check(model: onnxl.ModelProto, *_inputs) -> None:
    """Runs the two ORT-flatbuffer round-trips on *model*.

    Returning ``None`` tells :func:`make_test_class` that this is a
    model-level validator (no output comparison is performed).
    Backend-test inputs are ignored: this harness only checks the model
    serialization round-trip and does not run inference.

    Raises ``unittest.SkipTest`` while onnx-light has not implemented the
    ORT flatbuffer reader/writer yet (the C++ layer raises ``RuntimeError``
    from every ``SerializeFormat::kOrtFlatbuffers`` entry point).
    """
    assert isinstance(model, onnxl.ModelProto), f"Unexpected type {type(model)}"
    expected_bytes = _onnx_bytes(model)
    with tempfile.TemporaryDirectory() as tmp:
        light_path = os.path.join(tmp, "round_trip_light.ort")
        ort_path = os.path.join(tmp, "round_trip_ort.ort")

        # Step 1 + 2: save with onnx-light, reload with onnx-light, compare.
        try:
            _save_with_onnx_light_ort(model, light_path)
        except RuntimeError as exc:
            raise unittest.SkipTest(
                f"onnx-light ORT_FLATBUFFERS round-trip is not implemented yet: {exc}"
            ) from exc
        reloaded = _load_with_onnx_light_ort(light_path)
        assert (
            _onnx_bytes(reloaded) == expected_bytes
        ), "onnx-light → onnx-light ORT_FLATBUFFERS round-trip changed the model bytes"

        # Step 3 + 4: re-save through onnxruntime, reload with onnx-light, compare.
        _save_with_onnxruntime_ort(light_path, ort_path)
        reloaded_ort = _load_with_onnx_light_ort(ort_path)
        assert (
            _onnx_bytes(reloaded_ort) == expected_bytes
        ), "onnxruntime → onnx-light ORT_FLATBUFFERS round-trip changed the model bytes"


TestBackendRoundTripOrtFormat = make_test_class(ort_round_trip_check)


if __name__ == "__main__":
    unittest.main(verbosity=2)
