import os
import unittest
import carb
import omni.kit.test
import omni.ui as ui
from omni.kit import ui_test
from omni.kit.menu.utils import MenuItemDescription, MenuLayout, LayoutSourceSearch
from .utils import verify_menu_items, verify_menu_checked_items


class TestMenuLayoutUtils(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        pass

    async def tearDown(self):
        pass

    async def test_nested_submenu_layout(self):
        menu_placeholder = [MenuItemDescription(name="placeholder", show_fn=lambda: False)]

        menu_physics =  [ MenuItemDescription(name="Physics", sub_menu=[
                        MenuItemDescription(name="Debug"),
                        MenuItemDescription(name="Settings"),
                        MenuItemDescription(name="Demo Scenes"),
                        MenuItemDescription(name="Test Runner"),
                        MenuItemDescription(name="Character Controller")]) ]

        menu_blast =  [ MenuItemDescription(name="Blast", sub_menu=[
                        MenuItemDescription(name="Settings"),
                        MenuItemDescription(name="Documentation", sub_menu=[
                            MenuItemDescription(name="Kit UI"),
                            MenuItemDescription(name="Programming"),
                            MenuItemDescription(name="USD Schemas")]) ])]

        menu_flow =  [ MenuItemDescription(name="Flow", sub_menu=[
                        MenuItemDescription(name="Presets"),
                        MenuItemDescription(name="Monitor")]) ]

        # add menu
        omni.kit.menu.utils.add_menu_items(menu_placeholder, "Window", 90)
        omni.kit.menu.utils.add_menu_items(menu_physics, "Window")
        omni.kit.menu.utils.add_menu_items(menu_blast, "Window")
        omni.kit.menu.utils.add_menu_items(menu_flow, "Window")
        await ui_test.human_delay()

        # ----------------------------------------------------
        # nested layout test
        # ----------------------------------------------------
        menu_layout = [
            MenuLayout.Menu("Window", [
                MenuLayout.SubMenu("Simulation", [
                    MenuLayout.Group("Flow", [
                        MenuLayout.Item("Presets", source="Window/Flow/Presets"),
                        MenuLayout.Item("Monitor", source="Window/Flow/Monitor"),
                    ]),
                    MenuLayout.Group("Blast", [
                        MenuLayout.Item("Settings", source="Window/Blast/Settings"),
                        MenuLayout.SubMenu("Documentation", [
                            MenuLayout.Item("Kit UI", source="Window/Blast/Documentation/Kit UI"),
                            MenuLayout.Item("Programming", source="Window/Blast/Documentation/Programming"),
                            MenuLayout.Item("USD Schemas", source="Window/Blast/Documentation/USD Schemas"),
                        ]),
                    ]),
                    MenuLayout.Group("Physics", [
                        MenuLayout.Item("Demo Scenes"),
                        MenuLayout.Item("Settings", source="Window/Physics/Settings"),
                        MenuLayout.Item("Debug"),
                        MenuLayout.Item("Test Runner"),
                        MenuLayout.Item("Character Controller")
                    ]),
                ]),
            ])
        ]
        omni.kit.menu.utils.add_layout(menu_layout)

        # verify layout
        verify_menu_items(self, [(ui.Menu, "Window", True), (ui.Menu, "Simulation", True), (ui.MenuItem, "placeholder", False), (ui.Menu, "Physics", False), (ui.Menu, "Blast", False), (ui.Menu, "Flow", False), (ui.Separator, "Flow", True), (ui.MenuItem, "Presets", True), (ui.MenuItem, "Monitor", True), (ui.Separator, "Blast", True), (ui.MenuItem, "Settings", True), (ui.Menu, "Documentation", True), (ui.Separator, "Physics", True), (ui.MenuItem, "Demo Scenes", True), (ui.MenuItem, "Settings", True), (ui.MenuItem, "Debug", True), (ui.MenuItem, "Test Runner", True), (ui.MenuItem, "Character Controller", True), (ui.MenuItem, "Kit UI", True), (ui.MenuItem, "Programming", True), (ui.MenuItem, "USD Schemas", True), (ui.Menu, "Documentation", False)])

        # remove layout
        omni.kit.menu.utils.remove_layout(menu_layout)
        self.assertTrue(omni.kit.menu.utils.get_menu_layout() == [])

        # remove menu
        omni.kit.menu.utils.remove_menu_items(menu_physics, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_blast, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_flow, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_placeholder, "Window")
        self.assertTrue(omni.kit.menu.utils.get_merged_menus() == {})

    async def test_remapped_layout_checkbox(self):
        menu_placeholder = [MenuItemDescription(name="placeholder", show_fn=lambda: False)]

        async def refresh_menus():
            await ui_test.menu_click("Tools", human_delay_speed=4, show=True)
            await ui_test.menu_click("More Tools", human_delay_speed=4, show=True)
            await ui_test.menu_click("More Tools", human_delay_speed=4, show=False)

        example_window_is_ticked = False
        def get_example_window_is_ticked():
            nonlocal example_window_is_ticked
            return example_window_is_ticked

        def toggle_example_window_is_ticked():
            nonlocal example_window_is_ticked
            example_window_is_ticked = not example_window_is_ticked

        menu_entry1 = [ MenuItemDescription(name="Example Window", ticked=True,  ticked_fn=get_example_window_is_ticked, onclick_fn=toggle_example_window_is_ticked) ]
        menu_entry2 = [ MenuItemDescription(name="Best Window Ever", ticked=True,  ticked_fn=get_example_window_is_ticked, onclick_fn=toggle_example_window_is_ticked) ]

        # add menu
        omni.kit.menu.utils.add_menu_items(menu_placeholder, "Tools")
        omni.kit.menu.utils.add_menu_items(menu_placeholder, "More Tools")
        omni.kit.menu.utils.add_menu_items(menu_entry1, "Window")
        omni.kit.menu.utils.add_menu_items(menu_entry2, "Window")
        await ui_test.human_delay()

        # more the "Example Window" and "Best Window Ever" to differrent menus
        menu_layout = [
            MenuLayout.Menu("Tools", [
                MenuLayout.Item("Stage Window", source="Window/Example Window")
                ]),
            MenuLayout.Menu("More Tools", [
                MenuLayout.Item("Another Stage Window", source="Window/Best Window Ever")
                ]),
            MenuLayout.Menu("More Cheese", [
                MenuLayout.Item("Not A Menu Item", source="No Menu/Worst Window Ever")
                ])
            ]
        omni.kit.menu.utils.add_layout(menu_layout)
        await ui_test.human_delay()

        # update checked status (checked) - refreshing window should refresh "Tools" and "More Tools"
        example_window_is_ticked = True
        omni.kit.menu.utils.refresh_menu_items("Window")
        # open menu as items are not refreshed on refresh_menu_items
        await refresh_menus()
        # verify
        verify_menu_checked_items(self, [(ui.Menu, "Tools", True, False, False), (ui.Menu, "More Tools", True, False, False), (ui.Menu, "Window", False, False, False), (ui.MenuItem, "Stage Window", True, True, True), (ui.MenuItem, "Another Stage Window", True, True, True)])

        # update checked status (unchecked) - refreshing window should refresh "Tools" and "More Tools"
        example_window_is_ticked = False
        omni.kit.menu.utils.refresh_menu_items("Window")
        # open menu as items are not refreshed on refresh_menu_items
        await refresh_menus()
        # verify
        verify_menu_checked_items(self, [(ui.Menu, "Tools", True, False, False), (ui.Menu, "More Tools", True, False, False), (ui.Menu, "Window", False, False, False), (ui.MenuItem, "Stage Window", True, True, False), (ui.MenuItem, "Another Stage Window", True, True, False)])

        # remove layout
        omni.kit.menu.utils.remove_layout(menu_layout)
        self.assertTrue(omni.kit.menu.utils.get_menu_layout() == [])

        # remove menu
        omni.kit.menu.utils.remove_menu_items(menu_entry2, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_entry1, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_placeholder, "Tools")
        omni.kit.menu.utils.remove_menu_items(menu_placeholder, "More Tools")
        self.assertTrue(omni.kit.menu.utils.get_merged_menus() == {})

    async def test_layout_local_source(self):
        menu_placeholder = [MenuItemDescription(name="placeholder", show_fn=lambda: False)]

        menu_physics =  [ MenuItemDescription(name="Physics", sub_menu=[
                        MenuItemDescription(name="Debug"),
                        MenuItemDescription(name="Settings"),
                        MenuItemDescription(name="Demo Scenes"),
                        MenuItemDescription(name="Test Runner"),
                        MenuItemDescription(name="Character Controller")]) ]

        menu_blast =  [ MenuItemDescription(name="Blast", sub_menu=[
                        MenuItemDescription(name="Settings"),
                        MenuItemDescription(name="Documentation", sub_menu=[
                            MenuItemDescription(name="Kit UI"),
                            MenuItemDescription(name="Programming"),
                            MenuItemDescription(name="USD Schemas")]) ])]

        menu_flow =  [ MenuItemDescription(name="Flow", sub_menu=[
                        MenuItemDescription(name="Presets"),
                        MenuItemDescription(name="Monitor")]) ]

        # add menu
        omni.kit.menu.utils.add_menu_items(menu_placeholder, "Window", 90)
        omni.kit.menu.utils.add_menu_items(menu_physics, "Window")
        omni.kit.menu.utils.add_menu_items(menu_blast, "Window")
        omni.kit.menu.utils.add_menu_items(menu_flow, "Window")
        await ui_test.human_delay()

        # ----------------------------------------------------
        # nested layout test
        # ----------------------------------------------------
        menu_layout = [
            MenuLayout.Menu("Window", [
                MenuLayout.SubMenu("Simulation", [
                    MenuLayout.Group("Flow", [
                        MenuLayout.Item("Presets", source="Window/Flow/Presets"),
                        MenuLayout.Item("Monitor", source="Window/Flow/Monitor"),
                    ]),
                    MenuLayout.Group("Blast", [
                        MenuLayout.Item("Settings", source="Window/Blast/Settings"),
                        MenuLayout.SubMenu("Documentation", [
                            MenuLayout.Item("Kit UI", source="Window/Blast/Documentation/Kit UI"),
                            MenuLayout.Item("Programming", source="Window/Blast/Documentation/Programming"),
                            MenuLayout.Item("USD Schemas", source="Window/Blast/Documentation/USD Schemas"),
                        ]),
                    ]),
                    MenuLayout.Group("Physics", [
                        MenuLayout.Item("Demo Scenes"),
                        MenuLayout.Item("Settings", source="Window/Physics/Settings"),
                        MenuLayout.Item("Debug"),
                        MenuLayout.Item("Test Runner"),
                        MenuLayout.Item("Character Controller")
                    ]),
                ]),
                # these should not be moved as they are local only
                MenuLayout.Item("Demo Scenes", source_search=LayoutSourceSearch.LOCAL_ONLY),
                MenuLayout.Item("Settings", source_search=LayoutSourceSearch.LOCAL_ONLY),
                MenuLayout.Item("Debug", source_search=LayoutSourceSearch.LOCAL_ONLY),
                MenuLayout.Item("Test Runner", source_search=LayoutSourceSearch.LOCAL_ONLY),
                MenuLayout.Item("Character Controller", source_search=LayoutSourceSearch.LOCAL_ONLY),
            ])
        ]
        omni.kit.menu.utils.add_layout(menu_layout)

        # verify layout
        verify_menu_items(self, [(ui.Menu, "Window", True), (ui.Menu, "Simulation", True), (ui.MenuItem, "placeholder", False), (ui.Menu, "Physics", False), (ui.Menu, "Blast", False), (ui.Menu, "Flow", False), (ui.Separator, "Flow", True), (ui.MenuItem, "Presets", True), (ui.MenuItem, "Monitor", True), (ui.Separator, "Blast", True), (ui.MenuItem, "Settings", True), (ui.Menu, "Documentation", True), (ui.Separator, "Physics", True), (ui.MenuItem, "Demo Scenes", True), (ui.MenuItem, "Settings", True), (ui.MenuItem, "Debug", True), (ui.MenuItem, "Test Runner", True), (ui.MenuItem, "Character Controller", True), (ui.MenuItem, "Kit UI", True), (ui.MenuItem, "Programming", True), (ui.MenuItem, "USD Schemas", True), (ui.Menu, "Documentation", False)])

        # remove layout
        omni.kit.menu.utils.remove_layout(menu_layout)
        self.assertTrue(omni.kit.menu.utils.get_menu_layout() == [])

        # remove menu
        omni.kit.menu.utils.remove_menu_items(menu_physics, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_blast, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_flow, "Window")
        omni.kit.menu.utils.remove_menu_items(menu_placeholder, "Window")
        self.assertTrue(omni.kit.menu.utils.get_merged_menus() == {})
