import carb
import unittest
import omni.kit.test
from omni.kit import ui_test
from omni.kit.test_suite.helpers import wait_stage_loading


class TestUsdAPI(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        from omni.kit.test_suite.helpers import arrange_windows

        await arrange_windows("Stage", 200)
        await omni.usd.get_context().new_stage_async()
        self._prim_path = omni.kit.property.usd.tests.create_test_stage()

    async def tearDown(self):
        pass

    async def test_usd_api(self):
        from omni.kit.property.usd import PrimPathWidget

        # test PrimPathWidget.*_button_menu_entry
        button_menu_entry = PrimPathWidget.add_button_menu_entry(
            path="TestUsdAPI/Test Entry", glyph=None, name_fn=None, show_fn=None, enabled_fn=None, onclick_fn=None
        )
        self.assertTrue(button_menu_entry != None)
        button_menu_entries = PrimPathWidget.get_button_menu_entries()
        self.assertTrue(button_menu_entry in button_menu_entries)
        PrimPathWidget.remove_button_menu_entry(button_menu_entry)
        button_menu_entries = PrimPathWidget.get_button_menu_entries()
        self.assertFalse(button_menu_entry in button_menu_entries)

        # test PrimPathWidget.*_path_item
        path_func_updates = 0

        def my_path_func():
            nonlocal path_func_updates
            path_func_updates += 1

        # select prim so prim path widget is drawn
        usd_context = omni.usd.get_context()
        usd_context.get_selection().set_selected_prim_paths([self._prim_path], True)

        PrimPathWidget.add_path_item(my_path_func)
        path_items = PrimPathWidget.get_path_items()
        self.assertTrue(my_path_func in path_items)
        await ui_test.human_delay(10)
        self.assertTrue(path_func_updates == 1)
        PrimPathWidget.rebuild()
        await ui_test.human_delay(10)
        self.assertTrue(path_func_updates == 2)
        PrimPathWidget.remove_path_item(my_path_func)
        self.assertFalse(my_path_func in path_items)
        PrimPathWidget.rebuild()
        await ui_test.human_delay(10)
        self.assertTrue(path_func_updates == 2)

        usd_context.get_selection().set_selected_prim_paths([], True)
