"""Tests exercising bundles in OmniGraph"""
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit


# ======================================================================
class TestOmniGraphBundles(ogts.OmniGraphTestCase):
    async def test_get_prims_use_target_prims(self):
        controller = og.Controller()
        keys = og.Controller.Keys

        graph_path = "/World/Graph"
        (graph, (_, get_prims), (surface1, surface2, surface3), _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadPrims", "omni.graph.nodes.ReadPrims"),
                    ("GetPrims", "omni.graph.nodes.GetPrims"),
                ],
                keys.CREATE_PRIMS: [
                    ("/Surface1", "Cube"),
                    ("/Surface2", "Cube"),
                    ("/Surface3", "Cube"),
                ],
                keys.CONNECT: [
                    ("ReadPrims.outputs_primsBundle", "GetPrims.inputs:bundle"),
                ],
            },
        )

        stage = omni.usd.get_context().get_stage()
        context = graph.get_default_graph_context()
        bundle = context.get_output_bundle(get_prims, "outputs_bundle")

        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(graph_path + "/ReadPrims.inputs:prims"),
            targets=[surface1.GetPath(), surface2.GetPath(), surface3.GetPath()],
        )

        # get two prims
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(graph_path + "/GetPrims.inputs:prims"),
            targets=[surface1.GetPath(), surface2.GetPath()],
        )
        graph.evaluate()
        self.assertEqual(bundle.get_child_bundle_count(), 2)

        # switch back to default to get all
        omni.kit.commands.execute(
            "SetRelationshipTargets",
            relationship=stage.GetPropertyAtPath(graph_path + "/GetPrims.inputs:prims"),
            targets=[],
        )
        graph.evaluate()
        self.assertEqual(bundle.get_child_bundle_count(), 3)

    async def test_bundle_child_producer_consumer(self):
        controller = og.Controller()
        keys = og.Controller.Keys

        graph_path = "/World/Graph"
        (graph, (_, child_producer_node, child_consumer_node), (surface1, surface2), _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("ReadPrims", "omni.graph.nodes.ReadPrims"),
                    ("BundleChildProducer", "omni.graph.test.BundleChildProducerPy"),
                    ("BundleChildConsumer", "omni.graph.test.BundleChildConsumerPy"),
                ],
                keys.CREATE_PRIMS: [
                    ("/Surface1", "Cube"),
                    ("/Surface2", "Cube"),
                ],
                keys.CONNECT: [
                    ("ReadPrims.outputs_primsBundle", "BundleChildProducer.inputs:bundle"),
                    ("BundleChildProducer.outputs_bundle", "BundleChildConsumer.inputs:bundle"),
                ],
            },
        )

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

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(graph_path + "/ReadPrims.inputs:prims"),
            target=surface1.GetPath(),
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(graph_path + "/ReadPrims.inputs:prims"),
            target=surface2.GetPath(),
        )

        graph.evaluate()
        self.assertEqual(child_producer_node.get_attribute("outputs:numChildren").get(), 2)
        self.assertEqual(child_consumer_node.get_attribute("outputs:numChildren").get(), 1)
        self.assertEqual(child_consumer_node.get_attribute("outputs:numSurfaces").get(), 2)

        graph.evaluate()
        self.assertEqual(child_producer_node.get_attribute("outputs:numChildren").get(), 2)
        self.assertEqual(child_consumer_node.get_attribute("outputs:numChildren").get(), 1)
        self.assertEqual(child_consumer_node.get_attribute("outputs:numSurfaces").get(), 2)

    async def __test_input_bundle_validity(self, bundle_properties_name):
        """Test if connecting/disconnecting invalidates the database and reinitializes BundleContents for input bundle"""
        controller = og.Controller()
        keys = og.Controller.Keys

        # Disconnected input bundle should be invalid
        (graph, (_, bundle_properties), _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CREATE_NODES: [
                    ("BundleConstructor", "omni.graph.nodes.BundleConstructor"),
                    ("BundleProperties", f"omni.graph.test.{bundle_properties_name}"),
                ],
            },
        )
        graph.evaluate()
        self.assertFalse(bundle_properties.get_attribute("outputs:valid").get())

        # Connected input bundle should be valid
        (graph, _, _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CONNECT: [
                    ("BundleConstructor.outputs_bundle", "BundleProperties.inputs:bundle"),
                ],
            },
        )
        graph.evaluate()
        self.assertTrue(bundle_properties.get_attribute("outputs:valid").get())

    # ----------------------------------------------------------------------
    async def test_is_input_bundle_valid_cpp(self):
        """Test if connecting/disconnecting invalidates the database and reinitializes BundleContents for input bundle for C++"""
        await self.__test_input_bundle_validity("BundleProperties")

    # ----------------------------------------------------------------------
    async def test_is_input_bundle_valid_py(self):
        """Test if connecting/disconnecting invalidates the database and reinitializes BundleContents for input bundle for Python"""
        await self.__test_input_bundle_validity("BundlePropertiesPy")

    async def __test_producer_consumer_bundle_change_tracking(self, producer_name, consumer_name):
        controller = og.Controller()
        keys = og.Controller.Keys

        #    +--------------------+       +---------------------+
        #    |                    |       |                     |
        #    |  Bundle Producer   -------->   Bundle Consumer   |
        #    |                    |       |                     |
        #    +--------------------+       +---------------------+
        #        (make changes)         (recompute if input changed)

        (graph, (producer_node, consumer_node), _, _) = controller.edit(
            "/World/Graph",
            {
                keys.CREATE_NODES: [
                    ("BundleProducer", f"omni.graph.test.{producer_name}"),
                    ("BundleConsumer", f"omni.graph.test.{consumer_name}"),
                ],
                keys.CONNECT: [
                    ("BundleProducer.outputs_bundle", "BundleConsumer.inputs:bundle"),
                ],
            },
        )

        context = graph.get_context()
        producer_bundle = context.get_output_bundle(producer_node, "outputs_bundle")
        consumer_bundle = context.get_output_bundle(consumer_node, "outputs_bundle")

        # First evaluation is necessary.
        graph.evaluate()
        self.assertTrue(consumer_node.get_attribute("outputs:hasOutputBundleChanged").get())

        # Next evaluation producer did not create any data.
        graph.evaluate()
        self.assertFalse(consumer_node.get_attribute("outputs:hasOutputBundleChanged").get())

        # Create a child at the producer, then check if consumer recomputes.
        producer_bundle.create_child_bundle("surface0")
        graph.evaluate()
        self.assertTrue(consumer_node.get_attribute("outputs:hasOutputBundleChanged").get())
        self.assertEqual(consumer_bundle.get_child_bundle_count(), 1)

        # Next evaluation producer did not create any data.
        graph.evaluate()
        self.assertFalse(consumer_node.get_attribute("outputs:hasOutputBundleChanged").get())

    # ----------------------------------------------------------------------
    async def test_producer_consumer_bundle_change_tracking_cpp(self):
        await self.__test_producer_consumer_bundle_change_tracking("BundleProducer", "BundleConsumer")

    # ----------------------------------------------------------------------
    async def test_producer_consumer_bundle_change_tracking_py(self):
        await self.__test_producer_consumer_bundle_change_tracking("BundleProducerPy", "BundleConsumerPy")

    # ----------------------------------------------------------------------
    async def test_bundle_to_target_connection(self):
        """Backwards compatibility test to open a USD file where node with a bundle and target is connected.
        This test is to check if the connection from USD file is properly deserialized.
        """
        (result, error) = await ogts.load_test_file("TestBundleToTargetConnection.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)
        graph_path = "/World/PushGraph"
        controller = og.Controller()
        graph = controller.graph(graph_path)
        await controller.evaluate(graph)

        node = controller.node(f"{graph_path}/test_node_bundle_to_target_connection_01")
        attr = controller.attribute("inputs:target", node)
        connections = attr.get_upstream_connections()
        self.assertEqual(len(connections), 1)
        self.assertEqual(
            connections[0].get_path(), f"{graph_path}/test_node_bundle_to_target_connection/outputs_bundle"
        )
