import carb
import omni
import omni.kit.test
import omni.usd
import omni.client
import omni.kit.widget.stage
import omni.kit.app
import omni.kit.ui_test as ui_test

from omni.kit.test_suite.helpers import arrange_windows
from pxr import UsdGeom


class TestStageWidget(omni.kit.test.AsyncTestCase):

    async def setUp(self):
        self.app = omni.kit.app.get_app()
        await arrange_windows("Stage", 800, 600)

        self.usd_context = omni.usd.get_context()
        await self.usd_context.new_stage_async()
        self.stage = omni.usd.get_context().get_stage()

    async def tearDown(self):
        pass

    async def wait(self, frames=4):
        for i in range(frames):
            await self.app.next_update_async()

    async def test_visibility_toggle(self):
        stage_tree = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        prim = self.stage.DefinePrim("/prim0", "Xform")
        await self.wait()

        prim_visibility_widget = stage_tree.find("**/ToolButton[*].name=='visibility'")
        self.assertTrue(prim_visibility_widget)
        await prim_visibility_widget.click()
        await self.wait()
        self.assertEqual(UsdGeom.Imageable(prim).ComputeVisibility(), UsdGeom.Tokens.invisible)

        prim_visibility_widget = stage_tree.find("**/ToolButton[*].name=='visibility'")
        self.assertTrue(prim_visibility_widget)
        await prim_visibility_widget.click()
        await self.wait()
        self.assertEqual(UsdGeom.Imageable(prim).ComputeVisibility(), UsdGeom.Tokens.inherited)

    async def test_search_widget(self):
        stage_tree = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        self.assertTrue(stage_tree)
        stage_model = stage_tree.widget.model

        search_field = ui_test.find("Stage//Frame/**/StringField[*].name=='search'")
        self.assertTrue(search_field)

        prim1 = self.stage.DefinePrim("/prim", "Xform")
        prim2 = self.stage.DefinePrim("/prim/prim1", "Xform")
        prim3 = self.stage.DefinePrim("/prim/prim2", "Xform")
        prim4 = self.stage.DefinePrim("/other", "Xform")
        await self.wait()

        await search_field.input("prim")
        await ui_test.emulate_keyboard_press(carb.input.KeyboardInput.ENTER)
        await self.wait()
        all_children = stage_model.get_item_children(None)
        all_names = [child.name for child in all_children]
        search_field.widget.model.set_value("")
        await search_field.input("")
        await self.wait()
        await ui_test.emulate_keyboard_press(carb.input.KeyboardInput.ENTER)
        await self.wait()
        self.assertTrue(prim4.GetName() not in all_names)
