# source: https://github.com/onnx/onnx/blob/main/onnx/test/basic_test.py
from __future__ import annotations

import io
import os
import pathlib
import tempfile
import unittest
import numpy as np
import onnx_light.onnx as onnxl
import onnx_light.onnx.helper as oh
from onnx_light.onnx import defs, parser
from onnx_light.ext_test_case import ExtTestCase


def _simple_model() -> onnxl.ModelProto:
    """Creates a minimal ModelProto instance for use in tests."""
    model = onnxl.ModelProto()
    model.ir_version = defs.onnx_ir_version()
    model.producer_name = "onnx-test"
    model.graph.name = "test"
    return model


def _simple_tensor() -> onnxl.TensorProto:
    """Creates a minimal float TensorProto instance for use in tests."""
    return oh.make_tensor(
        name="test-tensor",
        data_type=onnxl.TensorProto.FLOAT,
        dims=(2, 3, 4),
        vals=[x + 0.5 for x in range(24)],
    )


class TestIOModel(ExtTestCase):
    """Tests for saving and loading ModelProto (protobuf format only)."""

    def test_load_model_when_input_is_bytes(self) -> None:
        proto = _simple_model()
        proto_bytes = proto.SerializeToString()
        loaded_proto = onnxl.load(proto_bytes)
        self.assertEqual(proto, loaded_proto)

    def test_save_and_load_model_when_input_is_file_name(self) -> None:
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = os.path.join(temp_dir, "model.onnx")
            onnxl.save(proto, model_path)
            loaded_proto = onnxl.load(model_path)
            self.assertEqual(proto, loaded_proto)

    def test_save_and_load_model_when_input_is_pathlike(self) -> None:
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = pathlib.Path(temp_dir, "model.onnx")
            onnxl.save(proto, model_path)
            loaded_proto = onnxl.load(model_path)
            self.assertEqual(proto, loaded_proto)

    def test_save_and_load_model_when_input_is_file_descriptor(self) -> None:
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = os.path.join(temp_dir, "model_fd.onnx")
            flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
            if hasattr(os, "O_BINARY"):
                flags |= os.O_BINARY
            fd = os.open(model_path, flags, 0o666)
            try:
                proto.SerializeToFileDescriptor(fd)
            finally:
                os.close(fd)
            loaded_proto = onnxl.load(model_path)
            self.assertEqual(proto, loaded_proto)

    def test_save_to_file_descriptor_and_load(self) -> None:
        """Tests that save_to_file_descriptor() writes a valid model."""
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = os.path.join(temp_dir, "model_fd2.onnx")
            flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
            if hasattr(os, "O_BINARY"):
                flags |= os.O_BINARY
            fd = os.open(model_path, flags, 0o666)
            try:
                onnxl.save_to_file_descriptor(proto, fd)
            finally:
                os.close(fd)
            loaded_proto = onnxl.load(model_path)
            self.assertEqual(proto, loaded_proto)

    def test_save_model_when_input_is_file_like_object(self) -> None:
        """Tests that save() accepts a file-like object with a write() method."""
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = os.path.join(temp_dir, "model_stream.onnx")
            with open(model_path, "wb") as handle:
                onnxl.save(proto, handle)
            loaded_proto = onnxl.load(model_path)
            self.assertEqual(proto, loaded_proto)

    def test_serialize_to_ostream_and_load(self) -> None:
        """Tests that SerializeToOstream() writes valid bytes to a file-like object."""
        proto = _simple_model()
        buf = io.BytesIO()
        proto.SerializeToOstream(buf)
        loaded_proto = onnxl.load(buf.getvalue())
        self.assertEqual(proto, loaded_proto)

    def test_serialize_to_ostream_capital_s_and_load(self) -> None:
        """Tests that SerializeToOStream() (capital S) works as an alias."""
        proto = _simple_model()
        buf = io.BytesIO()
        proto.SerializeToOStream(buf)
        loaded_proto = onnxl.load(buf.getvalue())
        self.assertEqual(proto, loaded_proto)

    def test_save_model_file_like_with_external_data_raises(self) -> None:
        """Tests that save() raises when a file-like object is used with external data."""
        proto = _simple_model()
        buf = io.BytesIO()
        with self.assertRaises(ValueError) as ctx:
            onnxl.save(proto, buf, save_as_external_data=True)
        self.assertIn("file-like object", str(ctx.exception))

    def test_load_model_with_format_protobuf(self) -> None:
        """Tests that load() accepts format='protobuf' parameter."""
        proto = _simple_model()
        proto_bytes = proto.SerializeToString()
        loaded_proto = onnxl.load(proto_bytes, format="protobuf")
        self.assertEqual(proto, loaded_proto)

    def test_load_model_with_format_none(self) -> None:
        """Tests that load() accepts format=None parameter."""
        proto = _simple_model()
        with tempfile.TemporaryDirectory() as temp_dir:
            model_path = os.path.join(temp_dir, "model.onnx")
            onnxl.save(proto, model_path)
            loaded_proto = onnxl.load(model_path, format=None)
            self.assertEqual(proto, loaded_proto)

    def test_load_model_with_unsupported_format_raises(self) -> None:
        """Tests that load() raises ValueError for unsupported formats."""
        proto = _simple_model()
        proto_bytes = proto.SerializeToString()
        with self.assertRaises(ValueError) as ctx:
            onnxl.load(proto_bytes, format="ort")
        self.assertIn("Unsupported format", str(ctx.exception))
        self.assertIn("onnx-light only supports", str(ctx.exception))


class TestIOTensor(ExtTestCase):
    """Tests for saving and loading TensorProto (protobuf format only)."""

    def test_load_tensor_when_input_is_bytes(self) -> None:
        proto = _simple_tensor()
        proto_bytes = proto.SerializeToString()
        loaded_proto = onnxl.TensorProto()
        loaded_proto.ParseFromString(proto_bytes)
        self.assertEqual(proto.SerializeToString(), loaded_proto.SerializeToString())

    @unittest.skip("TensorProto.SerializeToFile is not yet implemented in onnx_light")
    def test_save_and_load_tensor_when_input_is_file_name(self) -> None:
        proto = _simple_tensor()
        with tempfile.TemporaryDirectory() as temp_dir:
            tensor_path = os.path.join(temp_dir, "tensor.data")
            proto.SerializeToFile(tensor_path)
            loaded_proto = onnxl.TensorProto()
            loaded_proto.ParseFromFile(tensor_path)
            self.assertEqual(proto.SerializeToString(), loaded_proto.SerializeToString())

    @unittest.skip("TensorProto.SerializeToFile is not yet implemented in onnx_light")
    def test_save_and_load_tensor_when_input_is_pathlike(self) -> None:
        proto = _simple_tensor()
        with tempfile.TemporaryDirectory() as temp_dir:
            tensor_path = pathlib.Path(temp_dir, "tensor.data")
            proto.SerializeToFile(str(tensor_path))
            loaded_proto = onnxl.TensorProto()
            loaded_proto.ParseFromFile(str(tensor_path))
            self.assertEqual(proto.SerializeToString(), loaded_proto.SerializeToString())


class TestBasicFunctions(ExtTestCase):
    def test_protos_exist(self) -> None:
        _ = onnxl.AttributeProto
        _ = onnxl.NodeProto
        _ = onnxl.GraphProto
        _ = onnxl.ModelProto

    def test_version_exists(self) -> None:
        model = onnxl.ModelProto()
        # When we create it, ir_version should not be set.
        self.assertFalse(model.has_ir_version())
        # Set ir_version to the current IR version.
        model.ir_version = defs.onnx_ir_version()
        model_bytes = model.SerializeToString()
        model.ParseFromString(model_bytes)
        self.assertTrue(model.has_ir_version())
        # Check if the version is correct.
        self.assertEqual(model.ir_version, defs.onnx_ir_version())

    def test_proto_bytesize_matches_serialized_length(self) -> None:
        model = _simple_model()
        self.assertEqual(model.ByteSize(), len(model.SerializeToString()))

        tensor = _simple_tensor()
        self.assertEqual(tensor.ByteSize(), len(tensor.SerializeToString()))

    def test_model_and_graph_str(self) -> None:
        # Check that str() works without error and contains expected fields.
        model = _simple_model()
        model_str = str(model)
        self.assertIn("ir_version", model_str)
        self.assertIn("onnx-test", model_str)
        self.assertIn("test", model_str)

        text_model = """
           <
             ir_version: 10,
             opset_import: [ "" : 19]
           >
           agraph (float[N] X) => (float[N] C)
           <
             float[1] weight = {1}
           >
           {
              C = Cast<to=1>(X)
           }
        """
        model = parser.parse_model(text_model)
        model_str = str(model)
        self.assertIn("ir_version: 10", model_str)
        self.assertIn("agraph", model_str)

        graph_str = str(model.graph)
        self.assertIn("agraph", graph_str)
        self.assertIn("Cast", graph_str)
        self.assertIn("weight", graph_str)

    def test_function_str(self) -> None:
        text = """
            <
            ir_version: 9,
            opset_import: [ "" : 15, "custom_domain" : 1],
            producer_name: "FunctionProtoTest",
            producer_version: "1.0",
            model_version: 1,
            doc_string: "A test model for model local functions."
          >
         agraph (float[N] x) => (float[N] out)
         {
            out = custom_domain.Selu<alpha=2.0, gamma=3.0>(x)
         }
         <
         domain: "custom_domain",
         opset_import: [ "" : 15],
         doc_string: "Test function proto"
         >
           Selu
           <alpha: float=1.67326319217681884765625, gamma: float=1.05070102214813232421875>
           (X) => (C)
           {
               constant_alpha = Constant<value_float: float=@alpha>()
               constant_gamma = Constant<value_float: float=@gamma>()
               alpha_x = CastLike(constant_alpha, X)
               gamma_x = CastLike(constant_gamma, X)
               exp_x = Exp(X)
               alpha_x_exp_x = Mul(alpha_x, exp_x)
               alpha_x_exp_x_ = Sub(alpha_x_exp_x, alpha_x)
               neg = Mul(gamma_x, alpha_x_exp_x_)
               pos = Mul(gamma_x, X)
               _zero = Constant<value_float=0.0>()
               zero = CastLike(_zero, X)
               less_eq = LessOrEqual(X, zero)
               C = Where(less_eq, neg, pos)
           }
        """
        model = parser.parse_model(text)
        model_str = str(model)
        self.assertIn("FunctionProtoTest", model_str)

        self.assertEqual(len(model.functions), 1)
        function_str = str(model.functions[0])
        self.assertIn("Selu", function_str)
        self.assertIn("custom_domain", function_str)
        self.assertIn("Test function proto", function_str)

    def test_repeated_proto_field_survives_container_clear(self) -> None:
        # Regression test: a Python reference to an element of a repeated proto
        # field must keep the underlying C++ object alive even after the
        # container is cleared via __delitem__ / clear.
        import numpy as np

        graph = oh.make_graph(
            [
                oh.make_node("MatMul", ["X", "W"], ["XW"]),
                oh.make_node("Add", ["XW", "B"], ["Z"]),
                oh.make_node("Relu", ["Z"], ["Y"]),
            ],
            "shape_inference_demo",
            inputs=[oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N", 4])],
            outputs=[oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)],
            initializer=[
                oh.make_tensor(
                    "W",
                    onnxl.TensorProto.FLOAT,
                    [4, 3],
                    np.zeros((4, 3), dtype=np.float32).flatten(),
                ),
                oh.make_tensor("B", onnxl.TensorProto.FLOAT, [3], np.zeros(3, dtype=np.float32)),
            ],
        )
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=8)

        # Take references via __getitem__, iteration and add() then clear the
        # underlying container; each held reference must remain usable.
        node_via_getitem = model.graph.node[0]
        iter_nodes = list(model.graph.node)
        new_node = model.graph.node.add()
        new_node.op_type = "Identity"

        del model.graph.node[:]

        self.assertEqual(node_via_getitem.op_type, "MatMul")
        self.assertEqual([n.op_type for n in iter_nodes], ["MatMul", "Add", "Relu"])
        self.assertEqual(new_node.op_type, "Identity")

    def test_persistence(self):
        graph = oh.make_graph(
            [
                oh.make_node("MatMul", ["X", "W"], ["XW"]),
                oh.make_node("Add", ["XW", "B"], ["Z"]),
                oh.make_node("Relu", ["Z"], ["Y"]),
            ],
            "shape_inference_demo",
            inputs=[oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N", 4])],
            outputs=[oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, None)],
            initializer=[
                oh.make_tensor(
                    "W",
                    onnxl.TensorProto.FLOAT,
                    [4, 3],
                    np.zeros((4, 3), dtype=np.float32).flatten(),
                ),
                oh.make_tensor("B", onnxl.TensorProto.FLOAT, [3], np.zeros(3, dtype=np.float32)),
            ],
        )
        model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=8)
        node1 = model.graph.node[0]
        s1 = node1.SerializeToString()
        del model.graph.node[:]
        s2 = node1.SerializeToString()
        self.assertEqual(s1, s2)


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