"""Tests for OmniGraph Usd Listener """

import tempfile
from pathlib import Path

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


# ======================================================================
class TestOmniGraphUsdListener(ogts.OmniGraphTestCase):
    """Test for omni graph usd listener"""

    async def test_reparent_target_of_path_attribute(self):
        """Listener is required to update attributes with role set to path, when its target changes location"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, create_nodes, _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [("ConstantPath", "omni.graph.nodes.ConstantPath")],
                keys.CREATE_PRIMS: [("/Cube", "Cube"), ("/Xform", "Xform")],
            },
        )

        constant_path_node = create_nodes[0]
        value_attribute = constant_path_node.get_attribute("inputs:value")
        value_attribute.set("/Cube")

        # MovePrim command in this case will cause 'didAddNonInertPrim' and 'didRemoveNonInertPrim' flags to be
        # populated and reparenting takes place
        omni.kit.commands.execute("MovePrim", path_from="/Cube", path_to="/Xform/Cube")

        value = value_attribute.get()
        self.assertEqual(value, "/Xform/Cube")

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

            # save
            tmp_file_path = Path(tmp_dir_name) / "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
            (result, error) = await ogts.load_test_file(str(tmp_file_path))
            self.assertTrue(result, error)

        constant_path_node = og.get_node_by_path("/TestGraph/ConstantPath")
        value_attribute = constant_path_node.get_attribute("inputs:value")
        self.assertEqual(value_attribute.get(), "/Xform/Cube")

    async def test_rename_target_of_path_attribute(self):
        """Listener is required to update attributes with role set to path, when its target changes path"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (_, create_nodes, _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [("ConstantPath", "omni.graph.nodes.ConstantPath")],
                keys.CREATE_PRIMS: [("/Cube", "Cube")],
            },
        )

        constant_path_node = create_nodes[0]
        value_attribute = constant_path_node.get_attribute("inputs:value")
        value_attribute.set("/Cube")

        for src_dst in (("/Cube", "/Cube1"), ("/Cube1", "/Cube2"), ("/Cube2", "/Cube3")):
            omni.kit.commands.execute("MovePrim", path_from=src_dst[0], path_to=src_dst[1])
            self.assertEqual(value_attribute.get(), src_dst[1])

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

                # save
                tmp_file_path = Path(tmp_dir_name) / "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
                (result, error) = await ogts.load_test_file(str(tmp_file_path))
                self.assertTrue(result, error)

            constant_path_node = og.get_node_by_path("/TestGraph/ConstantPath")
            value_attribute = constant_path_node.get_attribute("inputs:value")
            self.assertEqual(value_attribute.get(), src_dst[1])

    async def test_extended_path_type(self):
        """Resolve extended type to path and confirm that has been subscribed to the listener"""

        class TestExtendedPathABIPy:
            @staticmethod
            def compute(db):
                return True

            @staticmethod
            def get_node_type() -> str:
                return "omni.graph.tests.TestExtendedPathABIPy"

            @staticmethod
            def initialize_type(node_type):
                node_type.add_extended_input(
                    "inputs:extended_input", "path,token", True, og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION
                )
                return True

        og.register_node_type(TestExtendedPathABIPy, 1)

        controller = og.Controller()
        keys = og.Controller.Keys
        (_, create_nodes, _, _) = controller.edit(
            "/TestGraph",
            {
                keys.CREATE_NODES: [("TestExtendedPath", "omni.graph.tests.TestExtendedPathABIPy")],
                keys.CREATE_PRIMS: [("/Cube", "Cube")],
            },
        )

        extended_node = create_nodes[0]
        extended_input = extended_node.get_attribute("inputs:extended_input")

        # set resolved type to path
        extended_input.set_resolved_type(og.Type(og.BaseDataType.UCHAR, 1, 1, og.AttributeRole.PATH))
        self.assertEqual(extended_input.get_resolved_type().role, og.AttributeRole.PATH)

        extended_input.set("/Cube")
        self.assertEqual(extended_input.get(), "/Cube")

        omni.kit.commands.execute("MovePrim", path_from="/Cube", path_to="/Cube1")

        value = extended_input.get()
        self.assertEqual(value, "/Cube1")
