"""Basic tests of the graph context"""

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.test
import omni.kit.undo
import omni.usd
import omni.usd.commands
from pxr import Usd


class TestOmniGraphContext(ogts.OmniGraphTestCase):
    """Tests graph context functionality"""

    TEST_GRAPH_PATH = "/TestGraph"

    # ----------------------------------------------------------------------
    async def test_graph_context_input_target_bundles(self):
        """
        Tests that the getting the input target bundles on a graph context works correctly
        """

        usd_context = omni.usd.get_context()

        stage: Usd.Stage = usd_context.get_stage()

        controller = og.Controller()

        cube_prim = ogts.create_cube(stage, "Cube", (1, 0, 0))
        cone_prim = ogts.create_cone(stage, "Cone", (0, 1, 0))
        xform_prim = stage.DefinePrim("/Xform", "Xform")

        keys = og.Controller.Keys

        (graph, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ReadCube", "omni.graph.nodes.ReadPrims"),
                    ("ReadCone", "omni.graph.nodes.ReadPrims"),
                    ("WritePrims", "omni.graph.nodes.WritePrimsV2"),
                ],
                keys.CONNECT: [
                    ("ReadCube.outputs_primsBundle", "WritePrims.inputs:primsBundle"),
                    ("ReadCone.outputs_primsBundle", "WritePrims.inputs:primsBundle"),
                ],
            },
        )

        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/ReadCube.inputs:prims"),
            target=cube_prim.GetPath(),
        )
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/ReadCone.inputs:prims"),
            target=cone_prim.GetPath(),
        )
        omni.kit.commands.execute(
            "AddRelationshipTarget",
            relationship=stage.GetPropertyAtPath(f"{self.TEST_GRAPH_PATH}/WritePrims.inputs:prims"),
            target=xform_prim.GetPath(),
        )

        await controller.evaluate(graph)

        graph_context = graph.get_default_graph_context()
        self.assertTrue(graph_context.is_valid())
        write_prims_node = graph.get_node(f"{self.TEST_GRAPH_PATH}/WritePrims")
        self.assertTrue(write_prims_node.is_valid())
        read_cube_node = graph.get_node(f"{self.TEST_GRAPH_PATH}/ReadCube")
        self.assertTrue(read_cube_node.is_valid())
        read_cone_node = graph.get_node(f"{self.TEST_GRAPH_PATH}/ReadCone")
        self.assertTrue(read_cone_node.is_valid())

        write_prims_targets = graph_context.get_input_target_bundles(write_prims_node, "inputs:primsBundle")

        self.assertEqual(len(write_prims_targets), 2)

        [mbib0, mbib1] = write_prims_targets

        self.assertEqual(mbib0.get_child_bundle_count(), 1)
        self.assertEqual(mbib1.get_child_bundle_count(), 1)

        [spib0] = mbib0.get_child_bundles()
        [spib1] = mbib1.get_child_bundles()

        source_prim_paths = [
            spib0.get_attribute_by_name("sourcePrimPath").get(),
            spib1.get_attribute_by_name("sourcePrimPath").get(),
        ]

        # sort the paths because get_input_target_bundles does not define an ordering
        source_prim_paths.sort()

        self.assertEqual(source_prim_paths[0], "/Cone")
        self.assertEqual(source_prim_paths[1], "/Cube")
