"""OmniGraph Graph Settings Tests"""
import os
import tempfile

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 OmniGraphSchema


class TestOmniGraphSchemaGraphSettings(ogts.OmniGraphTestCase):
    """Tests for OmniGraph Graph Settings with Schema Enabled"""

    # ----------------------------------------------------------------------
    async def test_compute_settings_save_and_load(self):
        """Tests that creating, saving and loading compute graph settings writes and reads the correct settings"""

        pipeline_stages = list(og.GraphPipelineStage.__members__.values())
        evaluation_modes = list(og.GraphEvaluationMode.__members__.values())
        backing_types = [
            og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
            og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_WITHOUT_HISTORY,
            og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
        ]
        setting_tuples = zip(pipeline_stages[0:3], evaluation_modes[0:3], backing_types[0:3])

        index = 0

        for (pipeline_stage, evaluation_mode, backing_type) in setting_tuples:
            (graph, _, _, _) = og.Controller.edit(
                {
                    "graph_path": f"/World/TestGraph{index}",
                    "fc_backing_type": backing_type,
                    "pipeline_stage": pipeline_stage,
                    "evaluation_mode": evaluation_mode,
                }
            )
            path = graph.get_path_to_graph()
            self.assertIsNotNone(graph)

            with tempfile.TemporaryDirectory() as tmpdirname:
                # save the file
                tmp_file_path = os.path.join(tmpdirname, f"tmp{index}.usda")
                result = omni.usd.get_context().save_as_stage(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(tmp_file_path)
                self.assertTrue(result, error)

            graph = omni.graph.core.get_graph_by_path(path)
            self.assertIsNotNone(graph, f"Failed to load graph at {path} from {tmp_file_path}")
            self.assertEquals(graph.evaluation_mode, evaluation_mode)
            self.assertEquals(graph.get_pipeline_stage(), pipeline_stage)
            self.assertEquals(graph.get_graph_backing_type(), backing_type)
            index = index + 1

            await omni.usd.get_context().new_stage_async()

    # ----------------------------------------------------------------------
    async def test_evaluation_mode_changed_from_usd(self):
        """Tests that changing evaluation mode from usd is applied to the underlying graph"""

        controller = og.Controller()
        keys = og.Controller.Keys
        stage = omni.usd.get_context().get_stage()
        graph_path = "/World/TestGraph"

        (graph, nodes, _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                ]
            },
        )

        self.assertEquals(graph.evaluation_mode, og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC)

        # automatic mode, computes because no references exist
        count = nodes[0].get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(count + 1, nodes[0].get_compute_count())

        prim = stage.GetPrimAtPath(graph_path)
        self.assertTrue(prim.IsA(OmniGraphSchema.OmniGraph))
        schema_prim = OmniGraphSchema.OmniGraph(prim)

        # change the evaluation to instanced
        schema_prim.GetEvaluationModeAttr().Set("Instanced")
        count = nodes[0].get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_INSTANCED, graph.evaluation_mode)
        self.assertEqual(count, nodes[0].get_compute_count())

    # ----------------------------------------------------------------------
    async def test_evaluator_type_changed_from_usd(self):
        """Tests that changing evaluation mode from usd is applied to the underlying graph"""

        stage = omni.usd.get_context().get_stage()
        graph_path = "/World/TestGraph"
        controller = og.Controller()
        (graph, (read_time, constant_node, _), _, _) = controller.edit(
            {"graph_path": graph_path, "evaluator_name": "dirty_push"},
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("ReadTime", "omni.graph.nodes.ReadTime"),
                    ("ConstantInt", "omni.graph.nodes.ConstantInt"),
                    ("Sub", "omni.graph.test.SubtractDoubleC"),
                ],
                og.Controller.Keys.CONNECT: [
                    ("ReadTime.outputs:timeSinceStart", "Sub.inputs:a"),
                    ("ReadTime.outputs:timeSinceStart", "Sub.inputs:b"),
                ],
                og.Controller.Keys.SET_VALUES: [("ConstantInt.inputs:value", 3)],
            },
        )

        # skip the initial compute
        await omni.kit.app.get_app().next_update_async()

        # validate the graph ticks as expected
        count_time = read_time.get_compute_count()
        count_print = constant_node.get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(count_time + 1, read_time.get_compute_count(), "Expected time to tick")
        self.assertEqual(count_print, constant_node.get_compute_count(), "Expected print not to tick")

        prim = stage.GetPrimAtPath(graph_path)
        self.assertTrue(prim.IsA(OmniGraphSchema.OmniGraph))
        schema_prim = OmniGraphSchema.OmniGraph(prim)

        # change the graph to push by changing the USD.
        schema_prim.GetEvaluatorTypeAttr().Set("push")
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(graph.get_evaluator_name(), "push")
        self.assertEqual(count_time + 2, read_time.get_compute_count(), "Expected time to tick")
        self.assertEqual(count_print + 1, constant_node.get_compute_count(), "Expected print to tick")

        # change the graph back
        schema_prim.GetEvaluatorTypeAttr().Set("dirty_push")
        await omni.kit.app.get_app().next_update_async()  # skip the intial frame
        count_time = read_time.get_compute_count()
        count_print = constant_node.get_compute_count()

        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(graph.get_evaluator_name(), "dirty_push")
        self.assertEqual(count_time + 1, read_time.get_compute_count(), "Expected time to tick")
        self.assertEqual(count_print, constant_node.get_compute_count(), "Expected print not to tick")

    # ---------------------------------------------------------------------
    async def test_fabric_backing_changed_from_usd(self):
        stage = omni.usd.get_context().get_stage()
        graph_path = "/World/TestGraph"
        (graph, _, _, _) = og.Controller.edit(
            {"graph_path": graph_path, "fc_backing_type": og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED}
        )

        self.assertEqual(og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED, graph.get_graph_backing_type())

        prim = stage.GetPrimAtPath(graph_path)
        self.assertTrue(prim.IsA(OmniGraphSchema.OmniGraph))
        schema_prim = OmniGraphSchema.OmniGraph(prim)

        schema_prim.GetFabricCacheBackingAttr().Set("StageWithoutHistory")
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_WITHOUT_HISTORY, graph.get_graph_backing_type())

    # ---------------------------------------------------------------------
    async def test_pipeline_stage_changed_from_usd(self):
        """Tests the changing the pipeline stage in usd changes the underlying pipeline stage"""
        controller = og.Controller()
        keys = og.Controller.Keys
        stage = omni.usd.get_context().get_stage()
        graph_path = "/World/TestGraph"

        (graph, nodes, _, _) = controller.edit(
            graph_path,
            {
                keys.CREATE_NODES: [
                    ("OnTick", "omni.graph.action.OnTick"),
                ]
            },
        )

        self.assertEqual(og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION, graph.get_pipeline_stage())

        # validate the node updates
        count = nodes[0].get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(count + 1, nodes[0].get_compute_count())

        prim = stage.GetPrimAtPath(graph_path)
        self.assertTrue(prim.IsA(OmniGraphSchema.OmniGraph))
        schema_prim = OmniGraphSchema.OmniGraph(prim)

        # change the pipeline stage to on demand
        schema_prim.GetPipelineStageAttr().Set("pipelineStageOnDemand")
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_ONDEMAND, graph.get_pipeline_stage())

        # validate the node did not update
        count = nodes[0].get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(count, nodes[0].get_compute_count())

        # change back to simulation
        schema_prim.GetPipelineStageAttr().Set("pipelineStageSimulation")
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_SIMULATION, graph.get_pipeline_stage())

        # validate the node updates
        count = nodes[0].get_compute_count()
        await omni.kit.app.get_app().next_update_async()
        self.assertEqual(count + 1, nodes[0].get_compute_count())
