"""Basic tests of the action graph"""
import omni.client
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from omni.kit.commands import execute
from pxr import Gf, Sdf, Vt


# ======================================================================
class TestActionGraphIntegration(ogts.OmniGraphTestCase):
    """Integration tests for Action Graph"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    async def setUp(self):
        """Set up  test environment, to be torn down when done"""
        self._var_index = 0
        await super().setUp()
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        # ----------------------------------------------------------------------

    async def test_path_attribute(self):
        """Check that AG properly handles USD prim paths being changed at runtime"""
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()

        (_, _, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                "create_nodes": [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.action.SetPrimActive"),
                    ("Query", "omni.graph.nodes.IsPrimActive"),
                    ("PrimPath", "omni.graph.nodes.ConstantPath"),
                    ("BoolNot", "omni.graph.nodes.BooleanNot"),
                    ("Simple", "omni.graph.tutorials.SimpleData"),
                ],
                "create_prims": [("Test", {}), ("Test/Child", {}), ("Dummy", {}), ("Dummy/Child", {})],
                "connect": [
                    ("OnTick.outputs:tick", "Set.inputs:execIn"),
                    ("Query.outputs:active", "BoolNot.inputs:valueIn"),
                    ("BoolNot.outputs:valueOut", "Set.inputs:active"),
                    ("PrimPath.inputs:value", "Simple.inputs:a_path"),
                    ("Simple.outputs:a_path", "Query.inputs:prim"),
                ],
                "set_values": [
                    ("Set.inputs:prim", "/Test/Child"),
                    ("PrimPath.inputs:value", "/Test"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        # Move the prim.
        execute("MovePrim", path_from="/Test", path_to="/Test1")
        test = stage.GetPrimAtPath("/Test1/Child")

        # Graph should still work. First run should deactivate the prim.
        await controller.evaluate()
        self.assertFalse(test.IsActive())

        # Second run should make it active again.
        await controller.evaluate()
        self.assertTrue(test.IsActive())

        # Make sure that if we change the value, it can reconnect correctly to the path.

        # First change to an unrelated target.
        controller.edit(
            self.TEST_GRAPH_PATH, {"set_values": [("Set.inputs:prim", "/Dummy"), ("PrimPath.inputs:value", "/Dummy")]}
        )

        # Disconnected form the prim => evaluation shouldn't change a thing.
        self.assertTrue(test.IsActive())
        await controller.evaluate()
        self.assertTrue(test.IsActive())

        # Now reset to proper name.
        controller.edit(
            self.TEST_GRAPH_PATH,
            {"set_values": [("Set.inputs:prim", "/Test1/Child"), ("PrimPath.inputs:value", "/Test1")]},
        )

        # Graph should work. First run should deactivate the prim.
        await controller.evaluate()
        self.assertFalse(test.IsActive())

        # Second run should make it active again.
        await controller.evaluate()
        self.assertTrue(test.IsActive())

        # Move again the prim.
        execute("MovePrim", path_from="/Test1", path_to="/Test2")
        test = stage.GetPrimAtPath("/Test2/Child")

        # Graph should still work. First run should deactivate the prim.
        await controller.evaluate()
        self.assertFalse(test.IsActive())

        # Second run should make it active again.
        await controller.evaluate()
        self.assertTrue(test.IsActive())

    # ----------------------------------------------------------------------
    async def test_path_on_property(self):
        """Another test to check AG properly handles USD prim paths being changed at runtime"""
        #           Query--------->BoolNot-----+
        #           ^                          |
        #           |                          v
        # Simple----+        OnObjectChange-> Set
        #      |                               ^
        #      +-------------------------------+
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        (_, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                "create_nodes": [
                    ("OnObjectChange", "omni.graph.action.OnObjectChange"),
                    ("Set", "omni.graph.action.SetPrimActive"),
                    ("Query", "omni.graph.nodes.IsPrimActive"),
                    ("BoolNot", "omni.graph.nodes.BooleanNot"),
                    ("Simple", "omni.graph.tutorials.SimpleData"),
                ],
                "create_prims": [("Test", {}), ("Test/Child", {"myfloat": ("float", 0)})],
                "connect": [
                    ("OnObjectChange.outputs:changed", "Set.inputs:execIn"),
                    ("Simple.outputs:a_path", "Set.inputs:prim"),
                    ("Simple.outputs:a_path", "Query.inputs:prim"),
                    ("Query.outputs:active", "BoolNot.inputs:valueIn"),
                    ("BoolNot.outputs:valueOut", "Set.inputs:active"),
                ],
                "set_values": [
                    ("Simple.inputs:a_path", "/Test"),
                    ("OnObjectChange.inputs:onlyPlayback", False),
                    ("OnObjectChange.inputs:path", "/Test/Child.myfloat"),
                ],
            },
        )

        test = prims[1]
        attr = test.GetAttribute("myfloat")

        # First evaluation should do nothing.
        self.assertTrue(test.IsActive())
        await controller.evaluate()
        self.assertTrue(test.IsActive())

        # Modify the trigger value, should activate the graph.
        attr.Set(1)
        self.assertTrue(test.IsActive())
        await controller.evaluate()
        self.assertFalse(test.IsActive())

        # Moving the parent prim should keep everything working.
        execute("MovePrim", path_from="/Test", path_to="/Test1")
        test = stage.GetPrimAtPath("/Test1/Child")
        attr = test.GetAttribute("myfloat")
        # => first, the move shouldn't have triggered a changed.
        self.assertFalse(test.IsActive())
        await controller.evaluate()
        self.assertFalse(test.IsActive())
        # => then, graph is executed again if value is changed.
        attr.Set(2)
        self.assertFalse(test.IsActive())
        await controller.evaluate()
        self.assertTrue(test.IsActive())

    # ----------------------------------------------------------------------
    async def test_writeprimattribute_scaler(self):
        """Exercise WritePrimAttribute for scalar value"""

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: ("/World/Float", {"myfloat": ("float", 0)}),
                keys.CONNECT: ("OnTick.outputs:tick", "Set.inputs:execIn"),
                keys.SET_VALUES: [
                    ("Set.inputs:name", "myfloat"),
                    ("Set.inputs:primPath", "/World/Float"),
                    ("Set.inputs:usePath", True),
                    ("Set.inputs:value", 42.0, "float"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("myfloat").Get(), 42.0)

        # Same test but with bundle input
        stage.GetPrimAtPath(f"{self.TEST_GRAPH_PATH}/Set").GetRelationship("inputs:prim").SetTargets(
            [Sdf.Path("/World/Float")]
        )

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.SET_VALUES: [
                    ("Set.inputs:usePath", False),
                    ("Set.inputs:value", 0.0, "float"),
                ]
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("myfloat").Get(), 0.0)

    # ----------------------------------------------------------------------
    async def test_writeprimattribute_noexist(self):
        """Test WritePrimAttribute can handle attribute that exists in USD but not in FC"""

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

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: ("/World/Cube", "Cube"),
                keys.CONNECT: ("OnTick.outputs:tick", "Set.inputs:execIn"),
                keys.SET_VALUES: [
                    ("Set.inputs:name", "visibility"),
                    ("Set.inputs:primPath", "/World/Cube"),
                    ("Set.inputs:usePath", True),
                    ("Set.inputs:value", "invisible", "token"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        await controller.evaluate(graph)
        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("visibility").Get(), "invisible")

    # ----------------------------------------------------------------------
    async def test_readwrite_primattribute_primvar(self):
        """Test WritePrimAttribute and ReadPrimAttribute can handle primvars"""

        # The idea is to Read the primvar that exists, but does not have a value. It should resolve to the correct type
        # of float[3][] because the primvar is defined. Then take that value and append a color and feed that into
        # WritePrimAttribute to write it to a different prim which requires authoring the primvar and copying the value.
        #

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

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_PRIMS: [("/World/Cube1", "Cube"), ("/World/Cube2", "Cube")]}
        )
        self.assertIsNone(prims[0].GetAttribute("primvars:displayColor").Get())

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                    ("Get", "omni.graph.nodes.ReadPrimAttribute"),
                    ("ArraySetIndex", "omni.graph.nodes.ArraySetIndex"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Set.inputs:execIn"),
                    ("Get.outputs:value", "ArraySetIndex.inputs:array"),
                    ("ArraySetIndex.outputs:array", "Set.inputs:value"),
                ],
                keys.SET_VALUES: [
                    ("Set.inputs:name", "primvars:displayColor"),
                    ("Set.inputs:primPath", "/World/Cube2"),
                    ("Set.inputs:usePath", True),
                    ("Get.inputs:name", "primvars:displayColor"),
                    ("Get.inputs:primPath", "/World/Cube1"),
                    ("Get.inputs:usePath", True),
                    ("ArraySetIndex.inputs:resizeToFit", True),
                    ("ArraySetIndex.inputs:index", 0),
                    ("ArraySetIndex.inputs:value", (1, 0, 0), "float[3]"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        await controller.evaluate(graph)
        await controller.evaluate(graph)
        self.assertEqual(list(prims[1].GetAttribute("primvars:displayColor").Get()), [(1, 0, 0)])

    # ----------------------------------------------------------------------
    async def _test_writeprimattribute_tuple(self):
        """Exercise WritePrimAttribute for a tuple value"""

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

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: ("/World/Double3", {"mydouble3": ("double[3]", Gf.Vec3d(0, 0, 0))}),
                keys.CONNECT: ("OnTick.outputs:tick", "Set.inputs:execIn"),
                keys.SET_VALUES: [
                    ("Set.inputs:name", "mydouble3"),
                    ("Set.inputs:primPath", "/World/Double3"),
                    ("Set.inputs:usePath", True),
                    ("Set.inputs:value", Gf.Vec3d(42.0, 42.0, 42.0), "double[3]"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("mydouble3").Get(), Gf.Vec3d(42.0, 42.0, 42.0))

    # ----------------------------------------------------------------------
    async def test_writeprimattribute_array(self):
        """Exercise WritePrimAttribute for an array value"""

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

        big_array = Vt.IntArray(1000, [0 for _ in range(1000)])

        (graph, _, prims, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Set", "omni.graph.nodes.WritePrimAttribute"),
                ],
                keys.CREATE_PRIMS: ("/World/IntArray", {"myintarray": ("int[]", [0xFF for _ in range(10)])}),
                keys.CONNECT: ("OnTick.outputs:tick", "Set.inputs:execIn"),
                keys.SET_VALUES: [
                    ("Set.inputs:name", "myintarray"),
                    ("Set.inputs:primPath", "/World/IntArray"),
                    ("Set.inputs:usePath", True),
                    ("Set.inputs:value", big_array, "int[]"),
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )

        await controller.evaluate(graph)
        self.assertEqual(prims[0].GetAttribute("myintarray").Get(), big_array)

    async def test_resync_handling(self):
        """Test that resyncing a path only resets OG state when expected"""

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        cube = stage.DefinePrim("/World/scope1/cube", "Cube")
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (on_impulse, _, counter, extra), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnImpulse", "omni.graph.action.OnImpulseEvent"),
                    ("TickN", "omni.graph.action.Countdown"),
                    ("Counter0", "omni.graph.action.Counter"),
                    ("OnImpulseExtra", "omni.graph.action.OnImpulseEvent"),
                ],
                keys.CONNECT: [
                    ("OnImpulse.outputs:execOut", "TickN.inputs:execIn"),
                    ("TickN.outputs:finished", "Counter0.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnImpulse.inputs:onlyPlayback", False),
                    ("TickN.inputs:duration", 3),
                    ("OnImpulse.state:enableImpulse", True),
                ],
            },
        )
        count_attrib = og.Controller.attribute("outputs:count", counter)
        # tick 0 (setup)
        await controller.evaluate(graph)
        # tick 1
        await controller.evaluate(graph)
        # Removing the prim will trigger resync handling, we will verify that this doesn't
        # interrupt the latent node
        stage.RemovePrim(cube.GetPath())
        # tick 2
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.Controller.get(count_attrib), 0)
        # tick 3 (last tick in duration)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(count_attrib), 0)
        # tick 5 (finished triggers counter)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(count_attrib), 1)

        # Now check that removing a node prim does reset the state
        og.Controller.set(og.Controller.attribute("state:enableImpulse", on_impulse), True)
        # tick 0 (setup)
        await controller.evaluate(graph)
        # tick 1
        await controller.evaluate(graph)
        # Removing the prim will trigger resync handling, we will verify that this doesn't
        # interrupt the latent node
        stage.RemovePrim(extra.get_prim_path())
        # tick ( nothing )
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.Controller.get(count_attrib), 1)
        # tick 3 (nothing)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(count_attrib), 1)
        # tick 5 (nothing)
        await controller.evaluate(graph)
        self.assertEqual(og.Controller.get(count_attrib), 1)
