"""Action Graph with Compounds"""

import carb.events
import omni.graph.core as og
import omni.graph.core._unstable as ogu
import omni.graph.core.tests as ogts
import omni.kit.app


# ======================================================================
class TestActionGraphCompounds(ogts.OmniGraphTestCase):
    """Tests action graph evaluator functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

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

    # ------------------------------------------------------------------------
    async def test_create_action_subgraph_with_command(self):
        """Validates the create compound subgraph command in an action graph"""

        # "Add" will only compute if the subgraph is executed as a push graph OR
        # is flattened into the parent graph
        #            ┌───────────────┐
        #            │Subgraph       │
        #  ┌─────┐   │    ┌─────┐    │
        #  │Const├───►────► Add ├───►├────┐
        #  └─────┘   │    └──▲──┘    │    │
        #            │       │       │    │
        #            │ ┌─────┴┐      │    │
        #            │ │Const2│      │    │
        #            │ └──────┘      │    │
        #            │               │    │
        #            └───────────────┘    │
        # ┌───────┐                      ┌▼────────────┐
        # │OnTick ├──────────────────────►WriteVariable│
        # └───────┘                      └─────────────┘
        controller = og.Controller()
        keys = controller.Keys
        (graph, (_, _, const_2, add, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Const", "omni.graph.nodes.ConstantInt"),
                    ("Const2", "omni.graph.nodes.ConstantInt"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("WriteVariable", "omni.graph.core.WriteVariable"),
                ],
                keys.CREATE_VARIABLES: [
                    ("int_var", og.Type(og.BaseDataType.INT), 0),
                ],
                keys.CONNECT: [
                    ("Const.inputs:value", "Add.inputs:a"),
                    ("Const2.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "WriteVariable.inputs:value"),
                    ("OnTick.outputs:tick", "WriteVariable.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                    ("WriteVariable.inputs:variableName", "int_var"),
                ],
            },
        )
        await og.Controller.evaluate(graph)
        self.assertEquals(graph.get_variables()[0].get(graph.get_default_graph_context()), 0)
        controller.edit(
            self.TEST_GRAPH_PATH, {keys.SET_VALUES: [("Const.inputs:value", 22), ("Const2.inputs:value", 20)]}
        )

        ogu.cmds.ReplaceWithCompoundSubgraph(
            nodes=[add.get_prim_path(), const_2.get_prim_path()], compound_name="Subgraph"
        )
        await og.Controller.evaluate(graph)
        self.assertEquals(graph.get_variables()[0].get(graph.get_default_graph_context()), 42.0)

    # ------------------------------------------------------------------------
    async def test_event_node_in_subgraph(self):
        """
        Test that a compute-on-request node inside a subgraph works
        """
        controller = og.Controller(update_usd=True)
        keys = controller.Keys

        (graph, (compound_node,), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    (
                        "CompoundOuter",
                        {
                            keys.CREATE_NODES: [
                                (
                                    "CompoundInner",
                                    {
                                        keys.CREATE_NODES: [
                                            ("OnEvent", "omni.graph.action.OnMessageBusEvent"),
                                            ("Counter", "omni.graph.action.Counter"),
                                        ],
                                        keys.CONNECT: [("OnEvent.outputs:execOut", "Counter.inputs:execIn")],
                                        keys.SET_VALUES: [
                                            ("OnEvent.inputs:onlyPlayback", False),
                                            ("OnEvent.inputs:eventName", "testEvent"),
                                        ],
                                    },
                                )
                            ]
                        },
                    ),
                ],
            },
        )
        # One compute for the first-time subscribe.
        await omni.kit.app.get_app().next_update_async()
        msg = carb.events.type_from_string("testEvent")

        counter_attr = og.Controller.attribute(
            f"{compound_node.get_prim_path()}/Subgraph/CompoundInner/Subgraph/Counter.outputs:count"
        )
        await og.Controller.evaluate(graph)
        self.assertEqual(counter_attr.get(), 0)
        omni.kit.app.get_app().get_message_bus_event_stream().push(msg)
        # Wait for one kit update to allow the event-push mechanism to trigger the node callback.
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter_attr.get(), 1)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter_attr.get(), 1)
        omni.kit.app.get_app().get_message_bus_event_stream().push(msg)
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(counter_attr.get(), 2)
