"""Tests reparenting and renaming workflows"""
from math import isclose, nan

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


# ======================================================================
class TestRenameAndReparent(ogts.OmniGraphTestCase):
    """Run tests that exercises the renaming and reparenting of prims involved in the graph"""

    async def verify_basic_setup(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        graph = omni.graph.core.get_all_graphs()[0]

        await og.Controller.evaluate(graph)

        # make sure the setup is working to start with:
        bundle_inspector_node = og.Controller.node("/World/graph/bundle_inspector")
        bundle_inspector_output_values_attr = bundle_inspector_node.get_attribute("outputs:values")
        output_value = og.Controller.get(bundle_inspector_output_values_attr)

        # FSD may insert other values in the output_value
        def to_float(f):
            try:
                return float(f)
            except ValueError:
                return nan

        # self.assertEqual(float(output_value[0]), 118.0, 0.01)
        self.assertTrue(any(isclose(to_float(f), 118.0) for f in output_value))

        # mutate the "size" attribute of the cube, which should get picked up by the bundle inspector down
        # the line
        cube_prim = stage.GetPrimAtPath("/World/Cube")
        self.assertIsNotNone(cube_prim)
        size_attr = cube_prim.GetAttribute("size")
        self.assertIsNotNone(size_attr)
        size_attr.Set(119.0)
        await omni.kit.app.get_app().next_update_async()

        output_value = og.Controller.get(bundle_inspector_output_values_attr)
        self.assertTrue(any(isclose(to_float(f), 119.0) for f in output_value))

    async def test_reparent_graph(self):
        """Test basic reparenting operations.  Here we have a graph that we first reparent under 1 xform then
        we create another xform and parent the first xform under that
        """
        (result, error) = await ogts.load_test_file("TestBundleGraphReparent.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # make sure the setup is working to start with:
        await self.verify_basic_setup()

        cube_prim = stage.GetPrimAtPath("/World/Cube")
        size_attr = cube_prim.GetAttribute("size")

        # create an xform, and parent our graph under it.
        parent_xform_path = omni.usd.get_stage_next_free_path(stage, "/World/" + "parent_xform1", True)
        _ = stage.DefinePrim(parent_xform_path, "Xform")

        omni.kit.commands.execute("MovePrim", path_from="/World/graph", path_to="/World/parent_xform1/graph")
        size_attr.Set(120.0)
        await omni.kit.app.get_app().next_update_async()

        # must refetch the node and attribute, as these have been destroyed by the move:
        bundle_inspector_node = og.Controller.node("/World/parent_xform1/graph/bundle_inspector")
        bundle_inspector_output_values_attr = bundle_inspector_node.get_attribute("outputs:values")
        output_value = og.Controller.get(bundle_inspector_output_values_attr)

        # make sure the graph is still working
        self.assertEqual(float(output_value[0]), 120.0, 0.01)

        # create another xform, and parent our previous parent xform under it.
        parent_xform_path2 = omni.usd.get_stage_next_free_path(stage, "/World/" + "parent_xform2", True)
        stage.DefinePrim(parent_xform_path2, "Xform")

        omni.kit.commands.execute(
            "MovePrim", path_from="/World/parent_xform1", path_to="/World/parent_xform2/parent_xform1"
        )
        size_attr.Set(121.0)
        await omni.kit.app.get_app().next_update_async()

        # must refetch the node and attribute, as these have been destroyed by the move:
        bundle_inspector_node = og.Controller.node("/World/parent_xform2/parent_xform1/graph/bundle_inspector")
        bundle_inspector_output_values_attr = bundle_inspector_node.get_attribute("outputs:values")
        output_value = og.Controller.get(bundle_inspector_output_values_attr)

        # make sure the graph is still working
        self.assertEqual(float(output_value[0]), 121.0, 0.01)

    async def test_simple_rename(self):
        (result, error) = await ogts.load_test_file("TestBundleGraphReparent.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # make sure the setup is working to start with:
        await self.verify_basic_setup()

        cube_prim = stage.GetPrimAtPath("/World/Cube")
        size_attr = cube_prim.GetAttribute("size")

        # try renaming the graph:
        omni.kit.commands.execute("MovePrim", path_from="/World/graph", path_to="/World/graph_renamed")
        size_attr.Set(120.0)
        await omni.kit.app.get_app().next_update_async()

        # must refetch the node and attribute, as these have been destroyed by the move:
        bundle_inspector_node = og.Controller.node("/World/graph_renamed/bundle_inspector")
        bundle_inspector_output_values_attr = bundle_inspector_node.get_attribute("outputs:values")
        output_value = og.Controller.get(bundle_inspector_output_values_attr)

        # make sure the graph is still working
        self.assertEqual(float(output_value[0]), 120.0, 0.01)

        # try renaming a node within the graph:
        omni.kit.commands.execute(
            "MovePrim",
            path_from="/World/graph_renamed/bundle_inspector",
            path_to="/World/graph_renamed/bundle_inspector_renamed",
        )
        size_attr.Set(121.0)
        await omni.kit.app.get_app().next_update_async()

        # must refetch the node and attribute, as these have been destroyed by the move:
        bundle_inspector_node = og.Controller.node("/World/graph_renamed/bundle_inspector_renamed")
        bundle_inspector_output_values_attr = bundle_inspector_node.get_attribute("outputs:values")
        output_value = og.Controller.get(bundle_inspector_output_values_attr)

        # make sure the graph is still working
        self.assertEqual(float(output_value[0]), 121.0, 0.01)

    async def test_reparent_primnode(self):
        """Test reparenting of prim nodes, which, because of all the special case handling needs its own test"""
        (result, error) = await ogts.load_test_file("TestBundleGraphReparent.usda", use_caller_subdirectory=True)
        self.assertTrue(result, error)

        controller = og.Controller()
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        # make sure the setup is working to start with:
        await self.verify_basic_setup()
        # verify bundle inspector is reading cube attribs
        count = og.Controller.get(controller.attribute("outputs:count", "/World/graph/bundle_inspector"))
        self.assertGreaterEqual(count, 9)
        # create an xform, and parent our graph under it.
        parent_xform_path = omni.usd.get_stage_next_free_path(stage, "/World/" + "parent_xform1", True)
        _ = stage.DefinePrim(parent_xform_path, "Xform")
        omni.kit.commands.execute("MovePrim", path_from="/World/Cube", path_to="/World/parent_xform1/Cube")
        await controller.evaluate()
        # check that bundle inspector is still functional
        count = og.Controller.get(controller.attribute("outputs:count", "/World/graph/bundle_inspector"))
        self.assertGreaterEqual(count, 9)
