""" Tests for omnigraph callbacks """

import numpy as np
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd

# Global variable for sharing callback location
_callback_path = None


# ======================================================================
class TestCallbacks(ogts.OmniGraphTestCase):

    _graph_path = "/World/TestGraph"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conn_count = 0
        self.disconn_count = 0
        self.attr_count = 0

    # ----------------------------------------------------------------------
    async def test_connection_callbacks(self):
        """
        Test attribute connect/disconnect callbacks.
        """
        # Create the graph.
        (_, (simple1, simple2), _, _) = og.Controller.edit(
            self._graph_path,
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("simple1", "omni.graph.tutorials.SimpleData"),
                    ("simple2", "omni.graph.tutorials.SimpleData"),
                ],
            },
        )

        self.conn_count = 0
        self.disconn_count = 0
        self.attr_count = 0

        def on_connection(attr1, attr2):
            self.conn_count += 1
            # Were the correct arguments passed?
            if (attr1.get_name() == "outputs:a_int") and (attr2.get_name() == "inputs:a_int"):
                self.attr_count += 1

        def on_disconnection(attr1, attr2):
            self.disconn_count += 1
            # Were the correct arguments passed?
            if (attr1.get_name() == "outputs:a_int") and (attr2.get_name() == "inputs:a_int"):
                self.attr_count += 1

        conn1_cb = simple1.register_on_connected_callback(on_connection)
        conn2_cb = simple2.register_on_connected_callback(on_connection)
        disconn1_cb = simple1.register_on_disconnected_callback(on_disconnection)

        self.assertEqual(self.conn_count, 0)
        self.assertEqual(self.disconn_count, 0)
        self.assertEqual(self.attr_count, 0)

        # Make a connection. We should get calls from both nodes.
        attr1 = simple1.get_attribute("outputs:a_int")
        attr2 = simple2.get_attribute("inputs:a_int")
        attr1.connect(attr2, True)

        self.assertEqual(self.conn_count, 2)
        self.assertEqual(self.disconn_count, 0)
        self.assertEqual(self.attr_count, 2)

        # Break the connection. We should only get a call from simple1.
        attr1.disconnect(attr2, True)

        self.assertEqual(self.conn_count, 2)
        self.assertEqual(self.disconn_count, 1)
        self.assertEqual(self.attr_count, 3)

        # Deregister one of the connection callbacks then make a connection.
        simple1.deregister_on_connected_callback(conn1_cb)
        attr1.connect(attr2, True)

        self.assertEqual(self.conn_count, 3)
        self.assertEqual(self.disconn_count, 1)
        self.assertEqual(self.attr_count, 4)

        # Deregistering the same callback again should have no effect.
        simple1.deregister_on_connected_callback(conn1_cb)

        # Cleanup.
        simple1.deregister_on_connected_callback(conn2_cb)
        simple1.deregister_on_connected_callback(disconn1_cb)

    # ----------------------------------------------------------------------
    async def test_path_changed_callback(self):
        """
        Test the ability to register callbacks on nodes so that they receive path change notices from USD
        """
        # Create the graph.
        (graph, (simple_node,), _, _) = og.Controller.edit(
            self._graph_path,
            {
                og.Controller.Keys.CREATE_NODES: [("simple", "omni.graph.tutorials.SimpleData")],
            },
        )
        k_want_print = False
        k_not_received = "not_received"
        global _callback_path
        _callback_path = k_not_received

        def on_path_changed_callback(changed_paths):
            global _callback_path
            _callback_path = changed_paths[-1]

        def _do_the_print(changed_paths, want_print: bool = k_want_print):
            return print(changed_paths) if k_want_print else None

        print_changed_paths = _do_the_print

        async def test_callback(should_work: bool):
            cb_handle = simple_node.register_on_path_changed_callback(on_path_changed_callback)
            print_changed_paths_cb_handle = simple_node.register_on_path_changed_callback(print_changed_paths)

            usd_context = omni.usd.get_context()
            stage = usd_context.get_stage()
            stage.DefinePrim("/World/xform1", "Xform")
            await og.Controller.evaluate(graph)
            self.assertEqual(_callback_path, "/World/xform1" if should_work else k_not_received)

            simple_node.deregister_on_path_changed_callback(cb_handle)
            simple_node.deregister_on_path_changed_callback(print_changed_paths_cb_handle)
            stage.DefinePrim("/World/xform2", "Xform")
            # This time we should not receive a callback, so ensure path is unchanged
            self.assertEqual(_callback_path, "/World/xform1" if should_work else k_not_received)
            await og.Controller.evaluate(graph)

        with og.Settings.temporary(og.Settings.ENABLE_PATH_CHANGED_CALLBACK, True):
            await test_callback(True)

        _callback_path = k_not_received

        with og.Settings.temporary(og.Settings.ENABLE_PATH_CHANGED_CALLBACK, False):
            await test_callback(False)

    # ----------------------------------------------------------------------
    async def test_register_value_changed_callback(self):
        """Test Attribute.register_value_changed callback"""

        connected_attr = None

        def on_value_changed(attr):
            nonlocal connected_attr
            connected_attr = attr

        controller = og.Controller()
        keys = og.Controller.Keys

        (_, (test_node_a, test_node_b, test_node_c), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("NodeWithAllTypes1", "omni.graph.test.TestAllDataTypes"),
                    ("NodeWithAllTypes2", "omni.graph.test.TestAllDataTypes"),
                    ("NodeWithAllTypes3", "omni.graph.test.TestAllDataTypes"),
                ]
            },
        )

        test_values = [
            ("inputs:a_bool", True),
            ("inputs:a_double", 42.0),
            ("inputs:a_float_2", (2.0, 3.0)),
            ("inputs:a_string", "test"),
        ]

        for name, value in test_values:
            src_attrib = test_node_b.get_attribute(name.replace("inputs:", "outputs:"))
            attrib = test_node_a.get_attribute(name)
            attrib.register_value_changed_callback(on_value_changed)

            # test the callback is hit when an attribute changes value
            og.Controller.set(attrib, value)
            self.assertEqual(connected_attr, attrib)

            # validate it is called on connection
            connected_attr = None
            og.Controller.connect(src_attrib, attrib)
            self.assertEqual(connected_attr, attrib)

            # validate it is called on disconnection
            connected_attr = None
            og.Controller.disconnect(src_attrib, attrib)
            self.assertEqual(connected_attr, attrib)

            # reset the callback, verify it is not called
            connected_attr = None
            attrib.register_value_changed_callback(None)
            og.Controller.set(attrib, value)
            self.assertIsNone(connected_attr)

        # Now validate it is hit when set with USD instead of OG

        read_attrib = None
        got_value = None

        def on_value_changed_2(_):
            nonlocal read_attrib
            nonlocal got_value
            # check if we have an attribute to read the value from
            if read_attrib:
                got_value = og.Controller.get(controller.attribute(read_attrib, test_node_c))

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        for name, value in test_values:
            attrib_c = test_node_c.get_attribute(name)
            attrib_c.register_value_changed_callback(on_value_changed_2)

            read_attrib = name
            got_value = None
            usd_attrib = stage.GetAttributeAtPath(attrib_c.get_path())
            usd_attrib.Set(value)
            if isinstance(got_value, np.ndarray):
                self.assertSequenceEqual(got_value.tolist(), value)
            else:
                self.assertEqual(got_value, value)
            attrib_c.register_value_changed_callback(None)

        # Verify callbacks are _not_ sent when disableInfoNoticeHandlingInPlayback=true
        with og.Settings.temporary(og.Settings.DISABLE_INFO_NOTICE_HANDLING_IN_PLAYBACK, True):
            for name, value in test_values:
                attrib_c = test_node_c.get_attribute(name)
                attrib_c.register_value_changed_callback(on_value_changed_2)

                read_attrib = name
                got_value = None
                usd_attrib = stage.GetAttributeAtPath(attrib_c.get_path())
                usd_attrib.Set(value)
                self.assertEqual(got_value, None)
                attrib_c.register_value_changed_callback(None)
