"""Tests relating to management of OmniGraph graphs"""
import tempfile
from pathlib import Path

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit
import omni.usd


class TestOmniGraphGraphManagement(ogts.OmniGraphTestCase):

    # --------------------------------------------------------------------------------------------------------------
    async def test_multiple_schema_root_graphs(self):
        """Check that a configuration with more than one root level graph can be created, saved, and loaded.
        This test creates three root-level graphs and then executes all of them to confirm that they are working
        properly. Each graph is a simple addition sequence, where one is an action graph to confirm that different
        graphs can use different evaluators.

        Const1 ---v
                  Add ---> SimpleData
        Const2 ---^

        OnTick --> Counter
        """
        keys = og.Controller.Keys

        (graph_1, (const_nodea_1, const_nodeb_1, _, simple_data_node_1), _, _) = og.Controller.edit(
            "/Graph1",
            {
                keys.CREATE_NODES: [
                    ("ConstA", "omni.graph.nodes.ConstantInt"),
                    ("ConstB", "omni.graph.nodes.ConstantInt"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("SimpleData", "omni.graph.tutorials.SimpleData"),
                ],
                keys.CONNECT: [
                    ("ConstA.inputs:value", "Add.inputs:a"),
                    ("ConstB.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "SimpleData.inputs:a_int"),
                ],
                keys.SET_VALUES: [
                    ("ConstA.inputs:value", 3),
                    ("ConstB.inputs:value", 7),
                ],
            },
        )
        # we'll be reading data directly in the graph: prevent any auto-instancing merging
        graph_1.set_auto_instancing_allowed(False)
        simple_data_node_1_path = simple_data_node_1.get_prim_path()
        graph_1_path = graph_1.get_path_to_graph()

        (graph_2, (const_nodea_2, const_nodeb_2, _, simple_data_node_2), _, _) = og.Controller.edit(
            "/Graph2",
            {
                keys.CREATE_NODES: [
                    ("ConstA", "omni.graph.nodes.ConstantInt"),
                    ("ConstB", "omni.graph.nodes.ConstantInt"),
                    ("Add", "omni.graph.nodes.Add"),
                    ("SimpleData", "omni.graph.tutorials.SimpleData"),
                ],
                keys.CONNECT: [
                    ("ConstA.inputs:value", "Add.inputs:a"),
                    ("ConstB.inputs:value", "Add.inputs:b"),
                    ("Add.outputs:sum", "SimpleData.inputs:a_int"),
                ],
                keys.SET_VALUES: [
                    ("ConstA.inputs:value", 33),
                    ("ConstB.inputs:value", 77),
                ],
            },
        )
        # we'll be reading data directly in the graph: prevent any auto-instancing merging
        graph_2.set_auto_instancing_allowed(False)
        simple_data_node_2_path = simple_data_node_2.get_prim_path()
        graph_2_path = graph_2.get_path_to_graph()

        (graph_3, (ontick_node, counter_node), _, _) = og.Controller.edit(
            {
                "graph_path": "/Graph3",
                "evaluator_name": "execution",
            },
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                    ("Counter", "omni.graph.action.Counter"),
                ],
                keys.CONNECT: [
                    ("OnTick.outputs:tick", "Counter.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnTick.inputs:onlyPlayback", False),
                ],
            },
        )
        # we'll be reading data directly in the graph: prevent any auto-instancing merging
        graph_3.set_auto_instancing_allowed(False)
        counter_node_path = counter_node.get_prim_path()
        graph_3_path = graph_3.get_path_to_graph()

        # ==============================================================================================================
        # OM-52190 specifies work that will prevent the need for this temporary setting. Without it, the values being
        #          set onto the attributes will not be part of USD, and hence will not be saved out to the file,
        #          causing the second half of the test to fail.
        stage = omni.usd.get_context().get_stage()
        stage.GetPrimAtPath(const_nodea_1.get_prim_path()).GetAttribute("inputs:value").Set(3)
        stage.GetPrimAtPath(const_nodeb_1.get_prim_path()).GetAttribute("inputs:value").Set(7)
        stage.GetPrimAtPath(const_nodea_2.get_prim_path()).GetAttribute("inputs:value").Set(33)
        stage.GetPrimAtPath(const_nodeb_2.get_prim_path()).GetAttribute("inputs:value").Set(77)
        stage.GetPrimAtPath(ontick_node.get_prim_path()).GetAttribute("inputs:onlyPlayback").Set(False)
        # ==============================================================================================================

        await og.Controller.evaluate([graph_1, graph_2, graph_3])

        self.assertEqual(og.Controller(("inputs:a_int", simple_data_node_1)).get(), 10)
        self.assertEqual(og.Controller(("inputs:a_int", simple_data_node_2)).get(), 110)
        self.assertTrue(og.Controller(("state:count", counter_node)).get() > 0)

        with tempfile.TemporaryDirectory() as tmpdirname:
            await omni.kit.app.get_app().next_update_async()

            # save the file
            tmp_file_path = Path(tmpdirname) / "tmp.usda"
            result = omni.usd.get_context().save_as_stage(str(tmp_file_path))
            self.assertTrue(result)

            # refresh the stage
            await omni.usd.get_context().new_stage_async()

            # reload the file back
            (result, error) = await ogts.load_test_file(str(tmp_file_path))
            self.assertTrue(result, error)

        graph_1 = og.get_graph_by_path(graph_1_path)
        graph_2 = og.get_graph_by_path(graph_2_path)
        graph_3 = og.get_graph_by_path(graph_3_path)
        simple_data_node_1 = og.get_node_by_path(simple_data_node_1_path)
        simple_data_node_2 = og.get_node_by_path(simple_data_node_2_path)
        counter_node = og.get_node_by_path(counter_node_path)

        self.assertIsNotNone(simple_data_node_1)
        self.assertIsNotNone(simple_data_node_2)
        self.assertIsNotNone(counter_node)

        await og.Controller.evaluate([graph_1, graph_2, graph_3])

        self.assertEqual(og.Controller(("inputs:a_int", simple_data_node_1)).get(), 10)
        self.assertEqual(og.Controller(("inputs:a_int", simple_data_node_2)).get(), 110)
        self.assertTrue(og.Controller(("state:count", counter_node)).get() > 0)

        # Close the stage to avoid dangling references to the graph. (OM-84680)
        await omni.usd.get_context().close_stage_async()
