import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.usd
from pxr import OmniGraphSchemaTools

GRAPH_TARGET_PRIM = og.INSTANCING_GRAPH_TARGET_PATH


class TestTargetType(ogts.OmniGraphTestCase):
    """Tests related to using the target attribute type"""

    _graph_path = "/World/TestGraph"

    # -------------------------------------------------------------------------
    async def test_target_path_substitution(self):
        """Test that the target accepts target path substitutions"""

        # Create the graph
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (prim_node,), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [("GetPrimsPaths", "omni.graph.nodes.GetPrimPaths")],
                keys.SET_VALUES: [("GetPrimsPaths.inputs:prims", [GRAPH_TARGET_PRIM])],
            },
        )
        await og.Controller.evaluate(graph)
        actual_result = controller.get(("outputs:primPaths", prim_node))
        self.assertEqual(self._graph_path, actual_result[0])

        # Change the graph path and verify the changes get picked up correctly
        controller.edit(
            self._graph_path, {keys.SET_VALUES: [("GetPrimsPaths.inputs:prims", f"{GRAPH_TARGET_PRIM}/Child")]}
        )
        await og.Controller.evaluate(graph)

        actual_result = controller.get(("outputs:primPaths", prim_node))
        self.assertEqual(f"{self._graph_path}/Child", actual_result[0])

    # -------------------------------------------------------------------------
    async def test_target_path_substitution_on_instances(self):
        """Tests that target paths substitution works as expected on instances"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (prim_node,), prims, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [("GetPrimsPaths", "omni.graph.nodes.GetPrimPaths")],
                keys.SET_VALUES: [("GetPrimsPaths.inputs:prims", GRAPH_TARGET_PRIM)],
                keys.CREATE_PRIMS: [
                    ("/World/Instance1", "Xform"),
                    ("/World/Instance2", "Xform"),
                    ("/World/Instance3", "Xform"),
                    ("/World/Instance4", "Xform"),
                    ("/World/Instance5", "Xform"),
                ],
            },
        )

        context = graph.get_default_graph_context()

        # Create the instances
        stage = omni.usd.get_context().get_stage()
        for p in prims:
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, p.GetPath(), self._graph_path)

        await og.Controller.evaluate(graph)

        for i in range(0, len(prims)):
            actual_result = controller.get(("outputs:primPaths", prim_node), instance=i)
            self.assertEqual(str(actual_result[0]), str(context.get_graph_target(i)))

        # Change the path and verify the changes propagate correctly
        controller.edit(
            self._graph_path, {keys.SET_VALUES: [("GetPrimsPaths.inputs:prims", f"{GRAPH_TARGET_PRIM}/Child")]}
        )
        await og.Controller.evaluate(graph)

        for i in range(0, len(prims)):
            actual_result = controller.get(("outputs:primPaths", prim_node), instance=i)
            self.assertEqual(str(actual_result[0]), f"{context.get_graph_target(i)}/Child")

    # -------------------------------------------------------------------------
    async def test_target_types_substitution_with_multiple_target_paths(self):
        """Test that the target accepts multiple target path substitutions"""

        # Create the graph
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (prim_node,), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [("GetPrimsPaths", "omni.graph.nodes.GetPrimPaths")],
                keys.SET_VALUES: [
                    ("GetPrimsPaths.inputs:prims", [GRAPH_TARGET_PRIM, "/World/Other", f"{GRAPH_TARGET_PRIM}/Child"])
                ],
            },
        )
        await og.Controller.evaluate(graph)
        actual_result = controller.get(("outputs:primPaths", prim_node))
        expected_result = [self._graph_path, "/World/Other", f"{self._graph_path}/Child"]
        self.assertEqual(expected_result, actual_result)

    # -------------------------------------------------------------------------
    async def test_target_types_substitution_in_connected_graphs(self):
        """Tests that the graph target replacements propagate through the graph"""

        # Create the graph
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_const_node, prim_node,), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("ConstantPrims", "omni.graph.nodes.ConstantPrims"),
                    ("GetPrimPaths", "omni.graph.nodes.GetPrimPaths"),
                ],
                keys.SET_VALUES: [("ConstantPrims.inputs:value", [GRAPH_TARGET_PRIM])],
                keys.CONNECT: [("ConstantPrims.inputs:value", "GetPrimPaths.inputs:prims")],
            },
        )

        # test the single evaluation case
        await og.Controller.evaluate(graph)
        actual_result = controller.get(("outputs:primPaths", prim_node))
        self.assertEqual(self._graph_path, actual_result[0])

        # create and instances and test the instanced case
        num_instances = 5
        stage = omni.usd.get_context().get_stage()
        for i in range(0, num_instances):
            p = stage.DefinePrim(f"/World/Instance{i}", "Xform")
            OmniGraphSchemaTools.applyOmniGraphAPI(stage, p.GetPath(), self._graph_path)

        await og.Controller.evaluate(graph)
        context = graph.get_default_graph_context()
        for i in range(0, num_instances):
            actual_result = controller.get(("outputs:primPaths", prim_node), instance=i)
            self.assertEqual(str(actual_result[0]), context.get_graph_target(i))

    # -------------------------------------------------------------------------
    async def test_target_types_substitution_in_referenced_graphs(self):

        # Create the graph
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, (_, _,), _, _) = controller.edit(
            self._graph_path,
            {
                keys.CREATE_NODES: [
                    ("ConstantPrims", "omni.graph.nodes.ConstantPrims"),
                    ("GetPrimPaths", "omni.graph.nodes.GetPrimPaths"),
                ],
                keys.SET_VALUES: [("ConstantPrims.inputs:value", [GRAPH_TARGET_PRIM])],
                keys.CONNECT: [("ConstantPrims.inputs:value", "GetPrimPaths.inputs:prims")],
            },
        )

        instance_path = "/World/Instance"

        stage = omni.usd.get_context().get_stage()
        prim = stage.DefinePrim(instance_path)
        prim.GetReferences().AddInternalReference(self._graph_path)

        await og.Controller.evaluate()

        instance_graph = og.Controller.graph(instance_path)
        self.assertTrue(instance_graph.is_valid())
        await og.Controller.evaluate(instance_graph)

        attr = og.Controller.attribute(f"{instance_path}/GetPrimPaths.outputs:primPaths")
        self.assertTrue(attr.is_valid())
        self.assertEqual([instance_path], og.Controller.get(attr))
