import sys
import re
import asyncio
import unittest
import carb
import omni.kit.test
import omni.ui as ui
from omni.kit import ui_test


class WindowStub():
    """stub window class for use with WidgetRef & """
    def undock(_):
        pass

    def focus(_):
        pass

window_stub = WindowStub()

class TestMenuAOVs(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        import omni.kit.material.library

        # 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()

    async def tearDown(self):
        pass

    def get_menubar(self) -> ui_test.WidgetRef:
        from omni.kit.mainwindow import get_main_window

        window = get_main_window()
        return ui_test.WidgetRef(widget=window._ui_main_window.main_menu_bar, path="MainWindow//Frame/MenuBar")

    def find_menu(self, menubar: ui_test.WidgetRef, path: str) -> ui_test.WidgetRef:
        for widget in menubar.find_all("**/"):
            if isinstance(widget.widget, ui.Menu) or isinstance(widget.widget, ui.MenuItem):
                if widget.widget.text.encode('ascii', 'ignore').decode().strip() == path:
                    return ui_test.WidgetRef(widget.widget, path="xxx", window=window_stub)
        return None

    async def test_aov_menus(self):
        from omni.kit.menu.aov import AOVMenuExtension

        # get root menu
        menubar = self.get_menubar()

        # new stage
        await omni.usd.get_context().new_stage_async()
        stage = omni.usd.get_context().get_stage()

        ## create prim
        omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")

        # verify one prim
        prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll() if prim.GetPath().pathString not in ['/Render', '/Render/Vars']]
        self.assertTrue(prim_list == ['/Sphere'])

        # use menu Create/AOV
        create_widget = self.find_menu(menubar, "Create")
        await create_widget.click(pos=create_widget.position+ui_test.Vec2(0, 5))
        aov_widget = self.find_menu(create_widget, "AOV")
        await aov_widget.click(pos=aov_widget.position+ui_test.Vec2(50, 5))

        # verify multiple prims
        prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll() if prim.GetPath().pathString not in ['/Render', '/Render/Vars']]
        self.assertFalse(prim_list == ['/Sphere'])

        # undo changes
        omni.kit.undo.undo()

        # verify one prim
        prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll() if prim.GetPath().pathString not in ['/Render', '/Render/Vars']]
        self.assertTrue(prim_list == ['/Sphere'])

        # new stage to prevent layout dialogs
        await omni.usd.get_context().new_stage_async()
