"""Test cases reproducing bugs that need to be fixed"""
import os
import tempfile

import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.test
import omni.usd
from omni.kit import commands as kit_commands


# ======================================================================
class TestBugs(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises graph functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    async def test_om_29829(self):
        """Undo of a connect raises a pxr.Tf.ErrorException"""

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

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SrcNode", "omni.graph.tutorials.ArrayData"),
                    ("DstNode", "omni.graph.tutorials.ArrayData"),
                ],
                keys.CONNECT: ("SrcNode.outputs:result", "DstNode.inputs:original"),
            },
        )

        (success, error) = kit_commands.execute("Undo")
        self.assertTrue(success, f"Failed to process an undo - {error}")

    # ----------------------------------------------------------------------
    async def test_crash_on_reload_referenced_og(self):
        """Tests OM-41534"""
        context = omni.usd.get_context()
        (result, error) = await ogts.load_test_file("ReferenceMaster.usd", use_caller_subdirectory=True)
        self.assertTrue(result, f"{error}")
        n_graphs = len(og.get_all_graphs())
        await context.reopen_stage_async()
        await omni.kit.app.get_app().next_update_async()
        # Verify we have the expected number of graphs
        graphs = og.get_all_graphs()
        self.assertEqual(len(graphs), n_graphs, f"{graphs}")

    # ----------------------------------------------------------------------
    async def test_get_node_type(self):
        """OM-41093: Test that node types returned from get_node_type match the actual underlying node type"""

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

        # One of each type of node is sufficient for the test as implementation is consistent for .ogn nodes
        python_node_type = "omni.graph.test.TupleArrays"
        cpp_node_type = "omni.graph.test.TypeResolution"

        (_, (python_node, cpp_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("PythonNode", python_node_type),
                    ("CppNode", cpp_node_type),
                ]
            },
        )

        self.assertEqual(python_node.get_node_type().get_node_type(), python_node_type)
        self.assertEqual(cpp_node.get_node_type().get_node_type(), cpp_node_type)

    # ----------------------------------------------------------------------
    async def test_duplicate_graph(self):
        """OM-41661: Tests that ComputeGraph can be duplicated"""
        (result, error) = await ogts.load_test_file("ReferencedGraph.usd", use_caller_subdirectory=True)
        context = omni.usd.get_context()
        stage = context.get_stage()

        self.assertTrue(result, f"{error}")
        graphs_list_1 = og.get_all_graphs()
        n_graphs = len(graphs_list_1)
        old_prim_path = "/World/PushGraph"
        old_prim = stage.GetPrimAtPath(old_prim_path)
        self.assertTrue(old_prim.IsValid())
        new_prim_path = omni.usd.get_stage_next_free_path(stage, old_prim_path, False)
        omni.kit.commands.execute("CopyPrim", path_from=old_prim_path, path_to=new_prim_path, exclusive_select=False)
        await omni.kit.app.get_app().next_update_async()
        new_prim = stage.GetPrimAtPath(new_prim_path)
        self.assertTrue(new_prim.IsValid())
        # Verify we have the expected number of graphs
        graphs_list_2 = og.get_all_graphs()
        self.assertGreater(len(graphs_list_2), n_graphs, f"{graphs_list_2}")

        # OM-43564: Temporarily comment out the test below
        # new_graph = next((g for g in graphs_list_2 if g.get_path_to_graph() == "/World/PushGraph_01"))
        # Verify that it is being evaluated
        # controller = og.Controller(og.Controller.attribute(f"{new_prim_path}/on_tick.outputs:timeSinceStart"))
        # t0 = controller.get()
        # await controller.evaluate(new_graph)
        # t1 = controller.get()
        # self.assertGreater(t1, t0)

    # ----------------------------------------------------------------------
    async def test_delete_graphs(self):
        """OM-42006: Tests that ComputeGraph can be deleted"""
        (result, error) = await ogts.load_test_file("ReferencedGraph.usd", use_caller_subdirectory=True)
        context = omni.usd.get_context()
        stage = context.get_stage()

        self.assertTrue(result, f"{error}")
        graphs_list_1 = og.get_all_graphs()
        n_graphs = len(graphs_list_1)
        src_prim_path = "/World/PushGraph"
        self.assertTrue(stage.GetPrimAtPath(src_prim_path).IsValid())
        new_prim_paths = []
        # Make a bunch of copies of our graph and then delete them
        for _ in range(10):
            new_prim_paths.append(omni.usd.get_stage_next_free_path(stage, src_prim_path, False))
            omni.kit.commands.execute(
                "CopyPrim", path_from=src_prim_path, path_to=new_prim_paths[-1], exclusive_select=False
            )
        await omni.kit.app.get_app().next_update_async()

        new_graphs = [og.get_graph_by_path(p) for p in new_prim_paths]
        for g in new_graphs:
            self.assertTrue(g)

        graphs_list_2 = og.get_all_graphs()
        self.assertEqual(len(graphs_list_2), n_graphs + len(new_prim_paths), f"{graphs_list_2}")
        # Delete all but the last 2 copies
        omni.kit.commands.execute("DeletePrims", paths=new_prim_paths[0:-2])
        # verify prims are gone and graphs have been invalidated
        for new_prim_path in new_prim_paths[0:-2]:
            self.assertFalse(stage.GetPrimAtPath(new_prim_path).IsValid())
        for g in new_graphs[0:-2]:
            self.assertFalse(g)

        # Verify we have the expected number of graphs
        graphs_list_3 = og.get_all_graphs()
        self.assertEqual(len(graphs_list_3), n_graphs + 2, f"{graphs_list_3}")

        # Now delete nodes from within the last 2
        nodes_to_delete = [p + "/write_prim_attribute" for p in new_prim_paths[-2:]] + [
            p + "/on_tick" for p in new_prim_paths[-2:]
        ]
        omni.kit.commands.execute("DeletePrims", paths=nodes_to_delete)

        # Check they are gone
        for p in nodes_to_delete:
            self.assertFalse(stage.GetPrimAtPath(p).IsValid())
        for p in new_prim_paths[-2:]:
            graph = og.get_graph_by_path(p)
            self.assertTrue(graph)
            self.assertFalse(graph.get_node("write_prim_attribute"))
            self.assertFalse(graph.get_node("on_tick"))

    # ----------------------------------------------------------------------
    async def test_om_43741(self):
        """Unresolved optional extended attribute fails to run compute"""

        keys = og.Controller.Keys
        controller = og.Controller()
        (graph, (test_node, _, _), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("TestNode", "omni.graph.test.TestOptionalExtendedAttribute"),
                    ("Input", "omni.graph.nodes.ConstantInt"),
                    ("Output", "omni.graph.tutorials.SimpleData"),
                ],
                keys.SET_VALUES: [
                    ("Input.inputs:value", 5),
                    ("Output.inputs:a_int", 11),
                ],
            },
        )
        await controller.evaluate(graph)

        result_other = og.Controller.attribute("outputs:other", test_node)
        result_optional = og.Controller.attribute("outputs:optional", test_node)

        # With nothing connected the node should run the compute on the "other" attributes
        self.assertEqual(3, result_other.get(), "No connections")
        with self.assertRaises(TypeError):
            result_optional.get()

        # With just the input connected neither output should be computed
        controller.edit(graph, {keys.CONNECT: ("Input.inputs:value", "TestNode.inputs:optional")})
        await controller.evaluate(graph)
        self.assertEqual(10, result_other.get(), "Only input connected")
        with self.assertRaises(TypeError):
            result_optional.get()

        # With input and output connected the optional output should be computed
        controller.edit(
            graph,
            {
                keys.CONNECT: ("TestNode.outputs:optional", "Output.inputs:a_int"),
            },
        )
        await controller.evaluate(graph)
        self.assertEqual(10, result_other.get(), "Both input and output connected")
        self.assertEqual(5, result_optional.get(), "Both input and output connected")

    # ----------------------------------------------------------------------
    async def test_delete_graphs_by_parent(self):
        """OM-45753: Tests that OmniGraphs are cleaned up when their parent is deleted"""
        context = omni.usd.get_context()
        stage = context.get_stage()
        cube_prim = stage.DefinePrim("/World/Cube", "Cube")
        cube2_prim = stage.DefinePrim("/World/Cube2", "Cube")
        graph1_path = "/World/Cube/TestGraph1"
        graph2_path = "/World/Cube/TestGraph2"

        keys = og.Controller.Keys
        controller1 = og.Controller()
        (graph1, nodes1, _, _) = controller1.edit(
            graph1_path,
            {
                keys.CREATE_NODES: [
                    ("Input", "omni.graph.nodes.ConstantInt"),
                ],
                keys.SET_VALUES: [
                    ("Input.inputs:value", 1),
                ],
            },
        )
        await controller1.evaluate(graph1)

        controller2 = og.Controller()
        (graph2, _, _, _) = controller2.edit(
            graph2_path,
            {
                keys.CREATE_NODES: [
                    ("Input", "omni.graph.nodes.ConstantInt"),
                ],
                keys.SET_VALUES: [
                    ("Input.inputs:value", 2),
                ],
            },
        )
        await controller2.evaluate(graph2)

        # verify we don't screw up when a sibling is deleted
        omni.kit.commands.execute("DeletePrims", paths=[cube2_prim.GetPath().pathString, nodes1[0].get_prim_path()])
        self.assertTrue(graph1)
        graphs = og.get_all_graphs()
        self.assertEqual(2, len(graphs))

        # delete the parent prim and verify OG graph is gone and our graph wrapper has been invalidated
        omni.kit.commands.execute("DeletePrims", paths=[cube_prim.GetPath().pathString])
        self.assertFalse(stage.GetPrimAtPath(graph1_path).IsValid())
        self.assertFalse(stage.GetPrimAtPath(graph2_path).IsValid())
        graphs = og.get_all_graphs()
        self.assertEqual(0, len(graphs))
        self.assertFalse(graph1)
        self.assertFalse(graph2)

    # ----------------------------------------------------------------------
    async def test_graphs_in_sublayer_load(self):
        """OM-46106:Tests the graphs located in added sublayers are correct loaded"""

        self.assertEqual(len(og.get_all_graphs_and_subgraphs()), 0)

        # insert a new layer with a single graph
        await ogts.insert_sublayer("TestGraphVariables.usda", True)
        self.assertEqual(len(og.get_all_graphs_and_subgraphs()), 1)

        # insert another layer containing multiple graphs
        await ogts.insert_sublayer("TestObjectIdentification.usda", True)
        self.assertEqual(len(og.get_all_graphs_and_subgraphs()), 7)

    # ----------------------------------------------------------------------
    async def test_om_70414(self):
        """
        Regression tests that loading a file with a non-zero start time doesn't
        crash on load.
        This can cause a simulation update without a stage attached.
        """
        (result, error) = await ogts.load_test_file("TestOM70414.usda", use_caller_subdirectory=True)
        self.assertTrue(result, f"{error}")

    # ----------------------------------------------------------------------
    async def test_om_73311(self):
        """Test that validates a resolved string using omni:graph:attrValue is accessable through the API"""

        (result, error) = await ogts.load_test_file("TestResolvedString.usda", use_caller_subdirectory=True)
        self.assertTrue(result, f"{error}")

        self.assertTrue(og.get_graph_by_path("/World/Graph").is_valid())
        self.assertTrue(og.Controller.node("/World/Graph/write_prim").is_valid())

        attribute = og.Controller.attribute("/World/Graph/write_prim.inputs:value")
        self.assertEquals(og.Controller.get(attribute), "chocolateChip")

    # ----------------------------------------------------------------------
    async def test_om_83905_token_array_crash(self):
        """Test that copying token array when restoring saved, unresolved values doesn't crash"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (node,), _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [("ToString", "omni.graph.nodes.ToString")],
                keys.CREATE_PRIMS: [("/Prim1", {}), ("/Prim2", {})],
                keys.SET_VALUES: [("ToString.inputs:value", {"type": "token[]", "value": ["/Prim1", "/Prim2"]})],
            },
        )
        await controller.evaluate(graph)

        attr = node.get_attribute("inputs:value")
        self.assertEqual(attr.get(), ["/Prim1", "/Prim2"])

        # Save and reload this file
        with tempfile.TemporaryDirectory() as tmp_dir_name:
            tmp_file_path = os.path.join(tmp_dir_name, "tmp.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
            (result, error) = await ogts.load_test_file(str(tmp_file_path))
            self.assertTrue(result, error)

            node = og.get_node_by_path("/TestGraph/ToString")
            attr = node.get_attribute("inputs:value")
            self.assertEqual(attr.get(), ["/Prim1", "/Prim2"])
