# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
"""Tests for the C++ ``onnx_op`` submodule exposed via nanobind."""

from __future__ import annotations

import unittest

from onnx_light.ext_test_case import ExtTestCase


class TestOnnxPyOp(ExtTestCase):
    def setUp(self) -> None:
        from onnx_light.onnx_py import _onnxpyprotoop  # type: ignore[attr-defined]

        self.mod = _onnxpyprotoop.onnx_op

    def test_tensor_type_and_to_type_string(self) -> None:
        self.assertEqual(self.mod.ToTypeString(self.mod.TensorType.kFloat), "tensor(float)")
        self.assertEqual(self.mod.ToTypeString(self.mod.TensorType.kInt64), "tensor(int64)")
        self.assertEqual(
            self.mod.ToTypeString(self.mod.TensorType.kSeqInt64), "seq(tensor(int64))"
        )

    def test_construct_formal_parameter(self) -> None:
        fp = self.mod.FormalParameter()
        fp.name = "X"
        fp.description = "input"
        fp.type = "T"
        self.assertEqual(fp.name, "X")
        self.assertEqual(fp.description, "input")
        self.assertEqual(fp.type, "T")

    def test_formal_parameter_repr(self) -> None:
        fp = self.mod.FormalParameter()
        fp.name = "X"
        fp.description = "input"
        fp.type = "T"
        self.assertEqual(repr(fp), "FormalParameter(name='X', type='T', description='input')")

    def test_construct_type_constraint_param(self) -> None:
        tc = self.mod.TypeConstraintParam()
        tc.type_param_str = "T"
        tc.allowed_type_strs = [self.mod.TensorType.kFloat, self.mod.TensorType.kDouble]
        tc.description = "any float"
        self.assertEqual(tc.type_param_str, "T")
        self.assertEqual(len(tc.allowed_type_strs), 2)
        self.assertEqual(tc.allowed_type_strs[0], self.mod.TensorType.kFloat)

    def test_type_constraint_param_repr(self) -> None:
        tc = self.mod.TypeConstraintParam()
        tc.type_param_str = "T"
        tc.allowed_type_strs = [self.mod.TensorType.kFloat, self.mod.TensorType.kDouble]
        tc.description = "any float"
        self.assertEqual(
            repr(tc),
            "TypeConstraintParam(type_param_str='T', "
            "allowed_type_strs=[TensorType.kFloat, TensorType.kDouble], "
            "description='any float')",
        )

    def test_construct_attribute_param_repr(self) -> None:
        attr = self.mod.AttributeParam()
        attr.name = "axis"
        attr.description = "axis for reduction"
        attr.type = self.mod.AttributeType.INT
        attr.required = False
        attr.default_value = 1
        self.assertEqual(
            repr(attr),
            "AttributeParam(name='axis', description='axis for reduction', "
            "type=AttributeType.INT, required=False, default_value='1')",
        )

    def test_construct_light_op_schema(self) -> None:
        fp = self.mod.FormalParameter()
        fp.name, fp.description, fp.type = "X", "in", "T"
        tc = self.mod.TypeConstraintParam()
        tc.type_param_str = "T"
        tc.allowed_type_strs = [self.mod.TensorType.kFloat]
        tc.description = "any float"
        schema = self.mod.LightOpSchema("MyOp", "ai.onnx", 1, "doc", [fp], [fp], [tc])
        self.assertEqual(schema.name, "MyOp")
        self.assertEqual(schema.domain, "ai.onnx")
        self.assertEqual(schema.since_version, 1)
        self.assertEqual(schema.doc, "doc")
        self.assertEqual(len(schema.inputs), 1)
        self.assertEqual(len(schema.outputs), 1)
        self.assertEqual(len(schema.type_constraints), 1)
        self.assertFalse(schema.has_function_implementation)
        # New fields: min_output/max_output default to len(outputs), deprecated defaults to False.
        self.assertEqual(schema.min_output, 1)
        self.assertEqual(schema.max_output, 1)
        self.assertFalse(schema.deprecated)
        # node_determinism defaults to Unknown (and is therefore not non_deterministic).
        self.assertEqual(schema.node_determinism, self.mod.LightOpSchema.NodeDeterminism.Unknown)
        self.assertFalse(schema.non_deterministic)
        schema.set_min_output(0).set_max_output(2**31 - 1).set_deprecated(True)
        self.assertEqual(schema.min_output, 0)
        self.assertEqual(schema.max_output, 2**31 - 1)
        self.assertTrue(schema.deprecated)
        schema.set_node_determinism(self.mod.LightOpSchema.NodeDeterminism.NonDeterministic)
        self.assertEqual(
            schema.node_determinism, self.mod.LightOpSchema.NodeDeterminism.NonDeterministic
        )
        self.assertTrue(schema.non_deterministic)

    def test_light_op_schema_repr(self) -> None:
        fp = self.mod.FormalParameter()
        fp.name, fp.description, fp.type = "X", "in", "T"
        tc = self.mod.TypeConstraintParam()
        tc.type_param_str = "T"
        tc.allowed_type_strs = [self.mod.TensorType.kFloat]
        tc.description = "any float"
        schema = self.mod.LightOpSchema("MyOp", "ai.onnx", 1, "doc", [fp], [fp], [tc])
        self.assertEqual(
            repr(schema),
            "LightOpSchema(name='MyOp', domain='ai.onnx', since_version=1, "
            "inputs=[FormalParameter(name='X', type='T', description='in')], "
            "outputs=[FormalParameter(name='X', type='T', description='in')], "
            "type_constraints=[TypeConstraintParam(type_param_str='T', "
            "allowed_type_strs=[TensorType.kFloat], description='any float')], "
            "attributes=[], has_function_implementation=False, "
            "min_output=1, max_output=1, deprecated=False, "
            "node_determinism=NodeDeterminism.Unknown)",
        )

    def test_get_all_onnx_op_schemas_with_history(self) -> None:
        schemas = self.mod.GetAllOnnxOpSchemasWithHistory()
        self.assertIsInstance(schemas, list)
        self.assertGreater(len(schemas), 0)
        for schema in schemas:
            self.assertIsInstance(schema, self.mod.LightOpSchema)
            self.assertTrue(schema.name)
            self.assertTrue(schema.domain)
            self.assertGreaterEqual(schema.since_version, 1)
        domains = {s.domain for s in schemas}
        self.assertIn("ai.onnx", domains)

    def test_get_all_onnx_op_schemas_with_history_op_type(self) -> None:
        all_schemas = self.mod.GetAllOnnxOpSchemasWithHistory()
        abs_schemas = self.mod.GetAllOnnxOpSchemasWithHistory("Abs")
        self.assertIsInstance(abs_schemas, list)
        self.assertGreater(len(abs_schemas), 0)
        for schema in abs_schemas:
            self.assertEqual(schema.name, "Abs")
        expected = sum(1 for s in all_schemas if s.name == "Abs")
        self.assertEqual(len(abs_schemas), expected)
        self.assertEqual(
            self.mod.GetAllOnnxOpSchemasWithHistory("ThisOpDoesNotExist", init_doc=False), []
        )

    def test_delayed_initializer_schema(self) -> None:
        delayed = self.mod.GetAllOnnxOpSchemasWithHistory("DelayedInitializer")
        self.assertEqual(len(delayed), 1)
        schema = delayed[0]
        self.assertEqual(schema.domain, "ai.rt")
        self.assertEqual(schema.since_version, 1)
        self.assertFalse(schema.has_function_implementation)
        self.assertEqual(schema.inputs, [])
        self.assertEqual(len(schema.outputs), 1)
        self.assertEqual(schema.outputs[0].name, "output")
        self.assertEqual(schema.outputs[0].type, "T")
        self.assertEqual(len(schema.type_constraints), 1)
        self.assertEqual(schema.type_constraints[0].type_param_str, "T")
        self.assertGreater(len(schema.type_constraints[0].allowed_type_strs), 10)
        self.assertNotIn("tensor(string)", schema.type_constraints[0].allowed_type_strs)
        self.assertEqual(
            [attribute.name for attribute in schema.attributes],
            ["shape", "dtype", "load_device", "runtime_device", "filename", "offset"],
        )
        self.assertTrue(all(attribute.required for attribute in schema.attributes))


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