import omni.kit.test
from re import I
from omni.ui.tests.test_base import OmniUiTest
from omni.kit.viewport.window import ViewportWindow
import omni.kit.ui_test as ui_test
from omni.kit.ui_test import Vec2
import omni.usd
import omni.kit.app
from pathlib import Path
import carb.input
import functools
from omni.kit.viewport.menubar.render import get_instance as _get_menubar_extension
from omni.kit.viewport.menubar.render import SingleRenderMenuItemBase as _SingleRenderMenuItemBase

CURRENT_PATH = Path(__file__).parent
TEST_DATA_PATH = CURRENT_PATH.parent.parent.parent.parent.parent.parent.joinpath("data").joinpath("tests")
TEST_WIDTH, TEST_HEIGHT = 600, 400


class TestSettingMenuWindow(OmniUiTest):
    async def setUp(self):
        self._golden_img_dir = TEST_DATA_PATH.absolute().joinpath("golden_img").absolute()

        await self.create_test_area(width=TEST_WIDTH, height=TEST_HEIGHT)

        self._vw = ViewportWindow('TestWindow', width=TEST_WIDTH, height=TEST_HEIGHT)
        self._viewport_api = self._vw.viewport_api

        await omni.kit.app.get_app().next_update_async()

    # After running each test
    async def tearDown(self):
        self._vw.destroy()
        del self._vw
        await super().tearDown()

    async def __show_render_menu(self, mouse_pos: Vec2 = None):
        # Enable mouse input
        app_window = omni.appwindow.get_default_app_window()
        for device in [carb.input.DeviceType.MOUSE]:
            app_window.set_input_blocking_state(device, None)

        await ui_test.emulate_mouse_move(mouse_pos if mouse_pos else Vec2(40, 40))
        await ui_test.emulate_mouse_click()
        await omni.kit.app.get_app().next_update_async()

    async def __close_render_menu(self):
        await ui_test.emulate_mouse_move(Vec2(TEST_WIDTH, TEST_HEIGHT))
        await ui_test.emulate_mouse_click()
        await self.wait_n_updates()

    async def test_render_menu_with_custom_type(self):
        """Test custom options area item. (This likely breaks order independent testing)"""
        menu_option_clicked = 0

        def _single_render_menu_item(*args, **kwargs):
            class SingleRenderMenuItem(_SingleRenderMenuItemBase):
                def _option_clicked(self):
                    nonlocal menu_option_clicked
                    menu_option_clicked += 1

            return SingleRenderMenuItem(*args, **kwargs)

        extension = _get_menubar_extension()
        self.assertIsNotNone(extension)
        try:
            await self.__show_render_menu()

            self.assertEqual(0, menu_option_clicked)

            # Simluate a click in options-area, but do it on Storm; not RTX
            await self.__show_render_menu(Vec2(205, 135))

            self.assertEqual(0, menu_option_clicked)

            extension.register_menu_item_type(
                functools.partial(_single_render_menu_item)
            )

            await self.wait_n_updates()
            await ui_test.emulate_mouse_click()
            await self.wait_n_updates()

            self.assertEqual(1, menu_option_clicked)
        finally:
            extension.register_menu_item_type(None)
            await self.__close_render_menu()
            await self.finalize_test_no_image()

    async def test_auto_manage_renderer_list(self):
        """Test auto management of the renderer menu based on extension load. (This likely breaks order independent testing)"""

        settings = carb.settings.get_settings()
        restore_value = settings.get("/renderer/enabled")
        try:
            async def change_available_renderers(golden_img_name: str, renderers: str = None):
                # Close the menu to avoid issues with golden images as menu shrink is animated over time
                if renderers:
                    settings.set("/renderer/enabled", renderers)
                await self.__show_render_menu()
                await self.capture_and_compare(
                    golden_img_dir=self._golden_img_dir, golden_img_name=golden_img_name
                )
                await self.__close_render_menu()


            await change_available_renderers("menubar_default_list_0.png")

            # Disable everything but RTX
            await change_available_renderers("menubar_rtx_only.png", "rtx")

            # Disable everything but Iray and Storm
            await change_available_renderers("menubar_pxr_iray_only.png", "pxr,iray")

            # Enable omni.hydra.index extension (Note it will enable RTX as well due to dependencies)
            omni.kit.app.get_app().get_extension_manager().set_extension_enabled_immediate("omni.hydra.index", True)
            await change_available_renderers("menubar_all_shipping.png")

            # Disbale omni.hydra.index extension, and list should be back to menubar_default_list
            omni.kit.app.get_app().get_extension_manager().set_extension_enabled_immediate("omni.hydra.index", False)
            await change_available_renderers("menubar_default_list_1.png")

        finally:
            settings.set("/renderer/enabled", restore_value if restore_value else "")
            await self.__close_render_menu()
            await self.finalize_test_no_image()
