import omni.kit.test
import asyncio
import carb
import omni.usd
import omni.ui as ui
from omni.kit import ui_test
from omni.kit.menu.utils import MenuItemDescription, MenuLayout


class TestMenuFile(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        # wait for material to be preloaded so create menu is complete & menus don't rebuild during tests
        await omni.kit.material.library.get_mdl_list_async()
        await ui_test.human_delay()

        self._future_test = None
        self._required_stage_event = -1
        self._stage_event_sub = omni.usd.get_context().get_stage_event_stream().create_subscription_to_pop(self._on_stage_event, name="omni.usd.menu.file")

    async def tearDown(self):
        pass

    def _on_stage_event(self, event):
        if self._future_test and int(self._required_stage_event) == int(event.type):
            self._future_test.set_result(event.type)

    async def reset_stage_event(self, stage_event):
        self._required_stage_event = stage_event
        self._future_test = asyncio.Future()

    async def wait_for_stage_event(self):
        async def wait_for_event():
            await self._future_test

        try:
            await asyncio.wait_for(wait_for_event(), timeout=30.0)
        except asyncio.TimeoutError:
            carb.log_error(f"wait_for_stage_event timeout waiting for {self._required_stage_event}")

        self._future_test = None
        self._required_stage_event = -1

    async def test_file_new_from_stage_template_empty(self):
        stage =  omni.usd.get_context().get_stage()
        layer_name =  stage.GetRootLayer().identifier if stage else "None"
        await self.reset_stage_event(omni.usd.StageEventType.OPENED)

        menu_widget = ui_test.get_menubar()
        await menu_widget.find_menu("File").click()
        await menu_widget.find_menu("New From Stage Template").click()

        # select empty and wait for stage open
        await menu_widget.find_menu("Empty").click()
        await self.wait_for_stage_event()

        # verify Empty stage
        stage =  omni.usd.get_context().get_stage()
        self.assertFalse(stage.GetRootLayer().identifier == layer_name)
        prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll()]
        self.assertTrue(prim_list == ['/World'])


    async def test_file_new_from_stage_template_sunlight(self):
        stage =  omni.usd.get_context().get_stage()
        layer_name =  stage.GetRootLayer().identifier if stage else "None"
        await self.reset_stage_event(omni.usd.StageEventType.OPENED)

        menu_widget = ui_test.get_menubar()
        await menu_widget.find_menu("File").click()
        await menu_widget.find_menu("New From Stage Template").click()

        # select Sunlight and wait for stage open
        await menu_widget.find_menu("Sunlight").click()
        await self.wait_for_stage_event()

        # verify Sunlight stage
        stage =  omni.usd.get_context().get_stage()
        self.assertFalse(stage.GetRootLayer().identifier == layer_name)
        prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll()]
        self.assertTrue(prim_list == ['/World', '/Environment', '/Environment/defaultLight'])
