"""Runs tests on the ABIs in the Python bindings"""
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import usdrt
from pxr import Sdf, Usd


# ==============================================================================================================
class TestAttributePythonBindings(ogts.OmniGraphTestCase):
    async def test_abi(self):
        """Tests that the bindings for omni.graph.core.Attribute operate as expected"""
        interface = og.Attribute
        ogts.validate_abi_interface(
            interface,
            instance_methods=[
                "connect",
                "connectEx",
                "connectPrim",
                "deprecation_message",
                "disconnect",
                "disconnectPrim",
                "get_all_metadata",
                "get_array",
                "get_attribute_data",
                "get_disable_dynamic_downstream_work",
                "get_downstream_connection_count",
                "get_downstream_connections_info",
                "get_downstream_connections",
                "get_extended_type",
                "get_handle",
                "get_metadata_count",
                "get_metadata",
                "get_name",
                "get_node",
                "get_path",
                "get_port_type",
                "get_resolved_type",
                "get_type_name",
                "get_union_types",
                "get_upstream_connection_count",
                "get_upstream_connections_info",
                "get_upstream_connections",
                "get",
                "is_array",
                "is_compatible",
                "is_connected",
                "is_deprecated",
                "is_dynamic",
                "is_runtime_constant",
                "is_valid",
                "register_value_changed_callback",
                "set_default",
                "set_disable_dynamic_downstream_work",
                "set_metadata",
                "set_resolved_type",
                "set",
                "update_attribute_value",
            ],
            static_methods=[
                "ensure_port_type_in_name",
                "get_port_type_from_name",
                "remove_port_type_from_name",
                "write_complete",
            ],
            properties=[
                "gpu_ptr_kind",
                "is_optional_for_compute",
            ],
            constants=[
                ("resolved_prefix", str),
            ],
        )

        # TODO: Now that we've confirmed the functions exist we also need to confirm that they work

    # --------------------------------------------------------------------------------------------------------------
    async def test_constants(self):
        """Test that the constants in the Attribute interface work as expected"""
        self.assertEqual(og.Attribute.resolved_prefix, "__resolved_")

    # --------------------------------------------------------------------------------------------------------------
    async def test_properties(self):
        """Test that the properties in the Attribute interface work as expected"""
        (_, (node,), _, _) = og.Controller.edit(
            "/TestGraph", {og.Controller.Keys.CREATE_NODES: ("TestNode", "omni.graph.test.PerturbPointsGpu")}
        )
        attr_points = og.Controller.attribute("inputs:points", node)
        self.assertEqual(attr_points.gpu_ptr_kind, og.PtrToPtrKind.NA)
        attr_points.gpu_ptr_kind = og.PtrToPtrKind.GPU
        self.assertEqual(attr_points.gpu_ptr_kind, og.PtrToPtrKind.GPU)

        self.assertFalse(attr_points.is_optional_for_compute)
        attr_points.is_optional_for_compute = True
        self.assertTrue(attr_points.is_optional_for_compute)

    # --------------------------------------------------------------------------------------------------------------
    async def test_static_methods(self):
        """Test that the static methods of the Attribute interface work as expected"""
        interface = og.Attribute
        p_in = og.AttributePortType.INPUT
        p_out = og.AttributePortType.OUTPUT
        p_state = og.AttributePortType.STATE

        # --------------------------------------------------------------------------------------------------------------
        # Test Attribute.ensure_port_type_in_name
        # Test data is a list of (expected_output, name_input, port_type, is_bundle)
        test_data = [
            ("inputs:attr", "inputs:attr", p_in, False),
            ("outputs:attr", "outputs:attr", p_out, False),
            ("state:attr", "state:attr", p_state, False),
            ("inputs:attr", "inputs:attr", p_in, True),
            ("outputs_attr", "outputs_attr", p_out, True),
            ("state_attr", "state_attr", p_state, True),
            ("inputs:attr", "attr", p_in, False),
            ("outputs:attr", "attr", p_out, False),
            ("state:attr", "attr", p_state, False),
            ("inputs:attr", "attr", p_in, True),
            ("outputs_attr", "attr", p_out, True),
            ("state_attr", "attr", p_state, True),
            ("inputs:attr:first", "inputs:attr:first", p_in, False),
            ("outputs:attr:first", "outputs:attr:first", p_out, False),
            ("state:attr:first", "state:attr:first", p_state, False),
            ("inputs:attr:first", "inputs:attr:first", p_in, True),
            ("outputs_attr_first", "outputs_attr_first", p_out, True),
            ("state_attr_first", "state_attr_first", p_state, True),
        ]

        for (expected, original, port_type, is_bundle) in test_data:
            msg = f"ensure_port_type_in_name({expected}, {original}, {port_type}, {is_bundle})"
            self.assertEqual(expected, interface.ensure_port_type_in_name(original, port_type, is_bundle), msg)

        # --------------------------------------------------------------------------------------------------------------
        # Test Attribute.get_port_type_from_name
        # Test data is a list of (expected_result, name_input)
        test_data = [
            (og.AttributePortType.INPUT, "inputs:attr"),
            (og.AttributePortType.OUTPUT, "outputs:attr"),
            (og.AttributePortType.STATE, "state:attr"),
            (og.AttributePortType.INPUT, "inputs:attr"),
            (og.AttributePortType.OUTPUT, "outputs_attr"),
            (og.AttributePortType.STATE, "state_attr"),
            (og.AttributePortType.UNKNOWN, "input:attr"),
        ]
        for (expected, original) in test_data:
            msg = f"get_port_type_from_name({expected}, {original})"
            self.assertEqual(expected, interface.get_port_type_from_name(original), msg)

        # --------------------------------------------------------------------------------------------------------------
        # Test Attribute.remove_port_type_from_name
        # Test data is a list of (expected_output, name_input, port_type, is_bundle)
        test_data = [
            ("attr", "inputs:attr", False),
            ("attr", "outputs:attr", False),
            ("attr", "state:attr", False),
            ("attr", "inputs:attr", True),
            ("attr", "outputs_attr", True),
            ("attr", "state_attr", True),
            ("attr", "attr", False),
            ("attr", "attr", False),
            ("attr", "attr", False),
            ("attr", "attr", True),
            ("attr", "attr", True),
            ("attr", "attr", True),
            ("attr:first", "inputs:attr:first", False),
            ("attr:first", "outputs:attr:first", False),
            ("attr:first", "state:attr:first", False),
            ("attr:first", "inputs:attr:first", True),
            ("attr_first", "outputs_attr_first", True),
            ("attr_first", "state_attr_first", True),
            ("input:attr", "input:attr", False),
            ("input:attr", "input:attr", True),
            ("output:attr", "output:attr", False),
            ("output_attr", "output_attr", True),
            ("attr", "attr", False),
            ("attr", "attr", True),
        ]

        for (expected, original, is_bundle) in test_data:
            msg = f"remove_port_type_from_name({expected}, {original}, {is_bundle})"
            self.assertEqual(expected, interface.remove_port_type_from_name(original, is_bundle), msg)

        # --------------------------------------------------------------------------------------------------------------
        # Test Attribute.write_complete
        # There's no easy way to judge that this worked as expected so instead just confirm that it returns okay
        (_, (node,), _, _) = og.Controller.edit(
            "/TestGraph", {og.Controller.Keys.CREATE_NODES: ("TestNode", "omni.graph.test.TestAllDataTypes")}
        )
        attrs = [
            og.Controller.attribute(name, node) for name in ["outputs:a_bool", "outputs:a_float", "outputs:a_double"]
        ]
        self.assertTrue(all(attr.is_valid() for attr in attrs))
        interface.write_complete(attrs)
        interface.write_complete(tuple(attrs))

    # --------------------------------------------------------------------------------------------------------------
    async def test_target_attribute_setting(self):
        """Tests various forms of setting an attribute with a target type"""

        (_, (node,), _, _) = og.Controller.edit(
            "/TestGraph",
            {
                og.Controller.Keys.CREATE_NODES: ("TestNode", "omni.graph.nodes.GetPrimPaths"),
                og.Controller.Keys.CREATE_PRIMS: [("/World/Prim1", "Xform"), ("/World/Prim2", "Xform")],
            },
        )
        attr = og.Controller.attribute("inputs:prims", node)

        # string
        og.Controller.set(attr, "/World/Prim1")
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim1")])

        # Sdf.Path
        og.Controller.set(attr, Sdf.Path("/World/Prim2"))
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim2")])

        # usdrt.Sdf.Path
        og.Controller.set(attr, usdrt.Sdf.Path("/World/Prim1"))
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim1")])

        # string list
        og.Controller.set(attr, ["/World/Prim1", "/World/Prim2"])
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim1"), usdrt.Sdf.Path("/World/Prim2")])

        # Sdf.Path list
        og.Controller.set(attr, [Sdf.Path("/World/Prim2"), Sdf.Path("/World/Prim1")])
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim2"), usdrt.Sdf.Path("/World/Prim1")])

        # usdrt.Sdf.Path list
        og.Controller.set(attr, [usdrt.Sdf.Path("/World/Prim1"), usdrt.Sdf.Path("/World/Prim2")])
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim1"), usdrt.Sdf.Path("/World/Prim2")])

        # empty list
        og.Controller.set(attr, [])
        self.assertEqual(og.Controller.get(attr), [])

        # mixed list
        og.Controller.set(attr, [Sdf.Path("/World/Prim2"), "/World/Prim1"])
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path("/World/Prim2"), usdrt.Sdf.Path("/World/Prim1")])

        # graph target path
        og.Controller.set(attr, og.INSTANCING_GRAPH_TARGET_PATH)
        self.assertEqual(og.Controller.get(attr), [usdrt.Sdf.Path(og.INSTANCING_GRAPH_TARGET_PATH)])

        # validate non-sdf paths throw
        omni.kit.commands.set_logging_enabled(False)
        with self.assertRaises(og.OmniGraphError):
            og.Controller.set(attr, 5)
        with self.assertRaises(og.OmniGraphError):
            og.Controller.set(attr, "!a non valid sdf path")
        with self.assertRaises(og.OmniGraphError):
            og.Controller.set(attr, ["/World/Prim1", 5])
        with self.assertRaises(og.OmniGraphError):
            og.Controller.set(attr, (usdrt.Sdf.Path("/World/Prim1"), usdrt.Sdf.Path("/World/Prim2")))
        omni.kit.commands.set_logging_enabled(True)

    # --------------------------------------------------------------------------------------------------------------
    async def test_objectlookup_attribute_with_paths(self):
        """Tests that og.Controller.attribute works with both strings and Sdf.Path objects for all types"""

        (_, (node,), _, _) = og.Controller.edit(
            "/TestGraph", {og.Controller.Keys.CREATE_NODES: ("TestNode", "omni.graph.test.TestAllDataTypes")}
        )

        for attr in node.get_attributes():
            if attr.get_port_type() not in (
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
            ):
                continue

            path = attr.get_path()
            self.assertEquals(attr, og.Controller.attribute(str(path)))
            self.assertEquals(attr, og.Controller.attribute(Sdf.Path(path)))

    # --------------------------------------------------------------------------------------------------------------
    async def test_objectlookup_usd_object_from_attributes(self):
        """Tests that og.ObjectLookup finds the correct usd object from an attribute"""
        controller = og.Controller(update_usd=True)
        keys = controller.Keys
        (_, (node,), _, _) = controller.edit(
            "/TestGraph", {keys.CREATE_NODES: ("TestNode", "omni.graph.test.TestAllDataTypes")}
        )

        for attr in node.get_attributes():
            port_type = attr.get_port_type()
            if port_type not in (
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
                og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT,
            ):
                continue

            # the exception is bundle outputs. They are neither attributes nor relationships.
            attribute_type = og.Controller.attribute_type(attr)
            if (
                port_type == og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT
                and attribute_type.role == og.AttributeRole.BUNDLE
            ):
                with self.assertRaises(og.OmniGraphError):
                    og.ObjectLookup.usd_attribute(attr)
                with self.assertRaises(og.OmniGraphError):
                    og.ObjectLookup.usd_relationship(attr)
                with self.assertRaises(og.OmniGraphError):
                    og.ObjectLookup.usd_property(attr)
                continue

            # everything else is an attribute or relationship
            path = attr.get_path()
            usd_property = og.ObjectLookup.usd_property(attr)
            self.assertTrue(usd_property.IsValid())
            self.assertEqual(usd_property, og.ObjectLookup.usd_property(Sdf.Path(path)))
            self.assertEqual(usd_property, og.ObjectLookup.usd_property(str(path)))

            # validate the property is a relationship or an attribute and that the corresponding inverse calls fail
            for item in [attr, Sdf.Path(path), str(path)]:
                if isinstance(usd_property, Usd.Relationship):
                    self.assertEqual(usd_property, og.ObjectLookup.usd_relationship(item))
                    with self.assertRaises(og.OmniGraphError):
                        og.ObjectLookup.usd_attribute(item)
                else:
                    self.assertEqual(usd_property, og.ObjectLookup.usd_attribute(item))
                    with self.assertRaises(og.OmniGraphError):
                        og.ObjectLookup.usd_relationship(item)
