# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import asyncio
import carb.settings
import omni.appwindow
import omni.kit.app
import omni.kit.test
import omni.timeline
import omni.kit.ui_test as ui_test
import omni.kit.widget.toolbar
import omni.kit.context_menu
import omni.kit.hotkeys.core
import omni.ui as ui
from carb.input import KeyboardInput as Key
from carb.input import KEYBOARD_MODIFIER_FLAG_SHIFT
from omni.kit.ui_test import Vec2
from omni.kit.widget.toolbar import Hotkey, SimpleToolButton, Toolbar, WidgetGroup, get_instance

from .helpers import reset_toolbar_settings

test_message_queue = []


class TestSimpleToolButton(SimpleToolButton):
    """
    Test of how to use SimpleToolButton
    """

    def __init__(self, icon_path):
        def on_toggled(c):
            test_message_queue.append(f"Test button toggled {c}")

        super().__init__(
            name="test_simple_tool_button",
            tooltip="Test Simple ToolButton",
            icon_path=f"{icon_path}/plus.svg",
            icon_checked_path=f"{icon_path}/plus.svg",
            hotkey=Key.U,
            toggled_fn=on_toggled,
            additional_style={"Button": { "color": 0xffffffff }}
        )


class TestToolButtonGroup(WidgetGroup):
    """
    Test of how to create two ToolButton in one WidgetGroup
    """

    def __init__(self, icon_path):
        super().__init__()
        self._icon_path = icon_path

    def clean(self):
        self._sub1 = None
        self._sub2 = None
        self._hotkey.clean()
        self._hotkey = None
        super().clean()

    def get_style(self):
        style = {
            "Button.Image::test1": {"image_url": f"{self._icon_path}/plus.svg"},
            "Button.Image::test1:checked": {"image_url": f"{self._icon_path}/minus.svg"},
            "Button.Image::test2": {"image_url": f"{self._icon_path}/minus.svg"},
            "Button.Image::test2:checked": {"image_url": f"{self._icon_path}/plus.svg"},
        }
        return style

    def create(self, default_size):
        def on_value_changed(index, model):
            if model.get_value_as_bool():
                self._acquire_toolbar_context()
            else:
                self._release_toolbar_context()
            test_message_queue.append(f"Group button {index} clicked")

        self._menu1 = omni.kit.context_menu.add_menu({ "name": "Test 1", "onclick_fn": None }, "test1", "omni.kit.widget.toolbar")
        self._menu2 = omni.kit.context_menu.add_menu({ "name": "Test 2", "onclick_fn": None }, "test1", "omni.kit.widget.toolbar")

        button1 = ui.ToolButton(
            name="test1",
            width=default_size,
            height=default_size,
            mouse_pressed_fn=lambda x, y, b, _: self._on_mouse_pressed(b, "test1"),
            mouse_released_fn=lambda x, y, b, _: self._on_mouse_released(b),
        )
        self._sub1 = button1.model.subscribe_value_changed_fn(lambda model, index=1: on_value_changed(index, model))

        button2 = ui.ToolButton(
            name="test2",
            width=default_size,
            height=default_size,
        )
        self._sub2 = button2.model.subscribe_value_changed_fn(lambda model, index=2: on_value_changed(index, model))

        self._hotkey = Hotkey(
            "toolbar::test2",
            Key.K,
            lambda: button2.model.set_value(not button2.model.get_value_as_bool()),
            lambda: self._is_in_context(),
        )

        # return a dictionary of name -> widget if you want to expose it to other widget_group
        return {"test1": button1, "test2": button2}


_MAIN_WINDOW_INSTANCE = None


class ToolbarApiTest(omni.kit.test.AsyncTestCase):
    WINDOW_NAME = "Main Toolbar Test"

    async def setUp(self):
        reset_toolbar_settings()

        self._icon_path = omni.kit.widget.toolbar.get_data_path().absolute().joinpath("icon").absolute()
        self._app = omni.kit.app.get_app()
        test_message_queue.clear()

        # If the instance doesn't exist, we need to create it for the test
        # Create a tmp window with the widget inside, Y axis because the tests are in the Y axis
        global _MAIN_WINDOW_INSTANCE
        if _MAIN_WINDOW_INSTANCE is None:
            _MAIN_WINDOW_INSTANCE = ui.ToolBar(
                self.WINDOW_NAME, noTabBar=False, padding_x=3, padding_y=3, margin=5, axis=ui.ToolBarAxis.Y
            )
            self._widget = get_instance()
            self._widget.set_axis(ui.ToolBarAxis.Y)
            self._widget.rebuild_toolbar(root_frame=_MAIN_WINDOW_INSTANCE.frame)
            await self._app.next_update_async()
            await self._app.next_update_async()

        self._main_dockspace = ui.Workspace.get_window("DockSpace")
        self._toolbar_handle = ui.Workspace.get_window(self.WINDOW_NAME)
        self._toolbar_handle.undock()
        await self._app.next_update_async()

        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.LEFT)
        await self._app.next_update_async()

    async def test_api(self):
        toolbar: Toolbar = omni.kit.widget.toolbar.get_instance()
        widget_simple = TestSimpleToolButton(self._icon_path)
        widget = TestToolButtonGroup(self._icon_path)

        toolbar.add_widget(widget, -100)
        toolbar.add_widget(widget_simple, -200)

        await self._app.next_update_async()

        self.assertIsNotNone(widget_simple.get_tool_button())

        # Check widgets are added and can be fetched
        self.assertIsNotNone(toolbar.get_widget("test_simple_tool_button"))
        self.assertIsNotNone(toolbar.get_widget("test1"))
        self.assertIsNotNone(toolbar.get_widget("test2"))

        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()

        manager = omni.kit.app.get_app().get_extension_manager()
        extension_name = omni.ext.get_extension_name(manager.get_extension_id_by_module(__name__))
        hotkey_registry = omni.kit.hotkeys.core.get_hotkey_registry()

        hotkey = self._get_hotkey('test_simple_tool_button::hotkey')
        self.assertIsNotNone(hotkey)

        hotkey_registry.edit_hotkey(hotkey, omni.kit.hotkeys.core.KeyCombination(Key.U, modifiers=KEYBOARD_MODIFIER_FLAG_SHIFT), None)

        tool_test_simple_pos = self._get_widget_center(toolbar, "test_simple_tool_button")
        tool_test1_pos = self._get_widget_center(toolbar, "test1")
        tool_test2_pos = self._get_widget_center(toolbar, "test2")

        # Test click on first simple button
        await self._emulate_click(tool_test_simple_pos)
        await self._app.next_update_async()

        # Test hot key on first simple button
        await self._emulate_keyboard_press(Key.U, KEYBOARD_MODIFIER_FLAG_SHIFT)
        await self._app.next_update_async()

        # Test click on 1/2 group button
        await self._emulate_click(tool_test1_pos)
        await self._app.next_update_async()

        # Test click on 2/2 group button
        await self._emulate_click(tool_test2_pos)
        await self._app.next_update_async()

        # Test click on 2/2 group button again
        await self._emulate_click(tool_test2_pos)
        await self._app.next_update_async()

        self.assertIsNone(ui.Menu.get_current())

        # Test right click on 1/2 group button
        await self._emulate_click(tool_test1_pos, right_click=True)
        await self._app.next_update_async()

        self.assertIsNotNone(ui.Menu.get_current())

        # Test left click on 1/2 group button to close the menu
        await self._emulate_click(tool_test1_pos, right_click=False)
        await self._app.next_update_async()

        self.assertIsNone(ui.Menu.get_current())

        # Test hold-left-click click on 1/2 group button
        await ui_test.input.emulate_mouse(ui_test.input.MouseEventType.MOVE, Vec2(*tool_test1_pos))
        await ui_test.input.emulate_mouse(ui_test.input.MouseEventType.LEFT_BUTTON_DOWN)
        await self._app.next_update_async()
        await asyncio.sleep(0.4)        
        await ui_test.input.emulate_mouse(ui_test.input.MouseEventType.LEFT_BUTTON_UP)
        await self._app.next_update_async()

        self.assertIsNotNone(ui.Menu.get_current())

        # Test left click on 1/2 group button to close the menu
        await self._emulate_click(tool_test1_pos)
        await self._app.next_update_async()

        self.assertIsNone(ui.Menu.get_current())

        # Making sure all button are triggered by checking the message queue
        expected_messages = [
            "Test button toggled True",
            "Test button toggled False",
            "Group button 1 clicked",
            "Group button 2 clicked",
            "Group button 2 clicked",
        ]

        self.assertEqual(expected_messages, test_message_queue)
        test_message_queue.clear()

        hotkey_registry.edit_hotkey(hotkey, omni.kit.hotkeys.core.KeyCombination(Key.U, modifiers=0), None)

        toolbar.remove_widget(widget)
        toolbar.remove_widget(widget_simple)

        widget.clean()
        widget_simple.clean()

        await self._app.next_update_async()

        # Check widgets are cleared
        self.assertIsNone(toolbar.get_widget("test_simple_tool_button"))
        self.assertIsNone(toolbar.get_widget("test1"))
        self.assertIsNone(toolbar.get_widget("test2"))


        # Change the Hotkey for the Play button
        play_hotkey = self._get_hotkey('toolbar::play')
        self.assertIsNotNone(play_hotkey)

        hotkey_registry.edit_hotkey(play_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.SPACE, modifiers=KEYBOARD_MODIFIER_FLAG_SHIFT), None)
        await self._app.next_update_async()
        tooltip_after_change = toolbar.get_widget('play').tooltip

        hotkey_registry.edit_hotkey(play_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.SPACE, modifiers=0), None)
        await self._app.next_update_async()
        tooltip_after_revert = toolbar.get_widget('play').tooltip

        self.assertEqual(tooltip_after_change, "Play (SHIFT + SPACE)")
        self.assertEqual(tooltip_after_revert, "Play (SPACE)")


        # Change the Hotkey for the Select Mode button
        select_mode_hotkey = self._get_hotkey('toolbar::select_mode')
        self.assertIsNotNone(select_mode_hotkey)

        hotkey_registry.edit_hotkey(select_mode_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.T, modifiers=KEYBOARD_MODIFIER_FLAG_SHIFT), None)
        await self._app.next_update_async()
        tooltip_after_change = toolbar.get_widget('select_mode').tooltip

        hotkey_registry.edit_hotkey(select_mode_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.T, modifiers=0), None)
        await self._app.next_update_async()
        tooltip_after_revert = toolbar.get_widget('select_mode').tooltip

        self.assertEqual(tooltip_after_change, "All Prim Types (SHIFT + T)")
        self.assertEqual(tooltip_after_revert, "All Prim Types (T)")


        # Change the Hotkey for the Move button
        move_op_hotkey = self._get_hotkey('toolbar::move')
        self.assertIsNotNone(move_op_hotkey)

        hotkey_registry.edit_hotkey(move_op_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.W, modifiers=KEYBOARD_MODIFIER_FLAG_SHIFT), None)
        await self._app.next_update_async()
        tooltip_after_change = toolbar.get_widget('move_op').tooltip

        hotkey_registry.edit_hotkey(move_op_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.W, modifiers=0), None)
        await self._app.next_update_async()
        tooltip_after_revert = toolbar.get_widget('move_op').tooltip

        from ..builtin_tools.transform_button_group import MOVE_TOOL_NAME
        self.assertEqual(tooltip_after_change, f"{MOVE_TOOL_NAME} (SHIFT + W)")
        self.assertEqual(tooltip_after_revert, f"{MOVE_TOOL_NAME} (W)")


        # Test the Select Mode hotkey button
        settings = carb.settings.get_settings()

        from ..builtin_tools.models.select_mode_model import SelectModeModel
        self.assertEqual(settings.get(SelectModeModel.PICKING_MODE_SETTING), SelectModeModel.PICKING_MODE_PRIMS)

        await self._emulate_keyboard_press(Key.T)
        await self._app.next_update_async()
        self.assertEqual(settings.get(SelectModeModel.PICKING_MODE_SETTING), SelectModeModel.PICKING_MODE_MODELS)

        await self._emulate_keyboard_press(Key.T)
        await self._app.next_update_async()
        self.assertEqual(settings.get(SelectModeModel.PICKING_MODE_SETTING), SelectModeModel.PICKING_MODE_PRIMS)

        # Test changing the Select hotkey
        select_hotkey = self._get_hotkey('toolbar::select')
        self.assertIsNotNone(select_hotkey)
        
        hotkey_registry.edit_hotkey(select_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.Q, modifiers=KEYBOARD_MODIFIER_FLAG_SHIFT), None)
        await self._app.next_update_async()

        hotkey_registry.edit_hotkey(select_hotkey, omni.kit.hotkeys.core.KeyCombination(Key.Q, modifiers=0), None)
        await self._app.next_update_async()

    async def test_context(self):
        toolbar = omni.kit.widget.toolbar.get_instance()
        widget_simple = TestSimpleToolButton(self._icon_path)
        widget = TestToolButtonGroup(self._icon_path)

        test_context = "test_context"

        toolbar.add_widget(widget, -100)
        toolbar.add_widget(widget_simple, -200, test_context)

        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()
        await self._app.next_update_async()

        tool_test_simple_pos = self._get_widget_center(toolbar, "test_simple_tool_button")
        tool_test1_pos = self._get_widget_center(toolbar, "test1")

        # Test click on first simple button
        # Add "Test button toggled True" to queue
        await self._emulate_click(tool_test_simple_pos)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), test_context)

        # Test hot key on 2/2 group button
        # It should have no effect since it's not "in context"
        # Add nothing to queue
        await self._emulate_keyboard_press(Key.K)
        await self._app.next_update_async()

        # Test click on first simple button again, should exit context
        # Add "Test button toggled False" to queue
        await self._emulate_click(tool_test_simple_pos)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), Toolbar.DEFAULT_CONTEXT)

        # Test hot key on 2/2 group button again
        # It should have effect because it's in default context
        # Add "Group button 2 clicked" to queue
        await self._emulate_keyboard_press(Key.K)
        await self._app.next_update_async()

        # Test click on first simple button again
        # Add "Test button toggled True" to queue
        await self._emulate_click(tool_test_simple_pos)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), test_context)

        # Test click on 1/2 group button so it takes context by force.
        # Add "Group button 1 clicked" to queue
        await self._emulate_click(tool_test1_pos)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), Toolbar.DEFAULT_CONTEXT)

        # Test hot key on first simple button, it should still work on default context.
        # Releasing an expired context.
        # Add "Test button toggled False" to queue
        await self._emulate_keyboard_press(Key.U)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), Toolbar.DEFAULT_CONTEXT)

        # Test hot key on first simple button, it should still work on default context.
        # Add "Test button toggled True" to queue
        await self._emulate_keyboard_press(Key.U)
        await self._app.next_update_async()

        self.assertEqual(toolbar.get_context(), test_context)

        expected_messages = [
            "Test button toggled True",
            "Test button toggled False",
            "Group button 2 clicked",
            "Test button toggled True",
            "Group button 1 clicked",
            "Test button toggled False",
            "Test button toggled True",
        ]

        self.assertEqual(expected_messages, test_message_queue)
        toolbar.remove_widget(widget)
        toolbar.remove_widget(widget_simple)

        widget.clean()
        widget_simple.clean()

        await self._app.next_update_async()

    def _get_widget_center(self, toolbar, id: str):
        return (
            toolbar.get_widget(id).screen_position_x + toolbar.get_widget(id).width / 2,
            toolbar.get_widget(id).screen_position_y + toolbar.get_widget(id).height / 2,
        )

    def _get_hotkey(self, action_id: str):
        manager = omni.kit.app.get_app().get_extension_manager()
        extension_name = omni.ext.get_extension_name(manager.get_extension_id_by_module(__name__))
        hotkey_registry = omni.kit.hotkeys.core.get_hotkey_registry()

        for hotkey in hotkey_registry.get_all_hotkeys_for_extension(extension_name):
            if hotkey.action_id == action_id:
                return hotkey

    async def _emulate_click(self, pos, right_click: bool = False):
        await ui_test.emulate_mouse_move_and_click(Vec2(*pos), right_click=right_click)

    async def _emulate_keyboard_press(self, key: Key, modifier: int = 0):
        await ui_test.emulate_keyboard_press(key, modifier)

    async def test_custom_select_types(self):
        toolbar: Toolbar = omni.kit.widget.toolbar.get_instance()

        entry_name = 'TestType'
        toolbar.add_custom_select_type(entry_name, ['test_type'])
        menu_dict = omni.kit.context_menu.get_menu_dict("select_mode", "omni.kit.widget.toolbar")
        self.assertIn(entry_name, [item['name'] for item in menu_dict])

        toolbar.remove_custom_select(entry_name)
        menu_dict = omni.kit.context_menu.get_menu_dict("select_mode", "omni.kit.widget.toolbar")
        self.assertNotIn(entry_name, [item['name'] for item in menu_dict])

    async def test_toolbar_methods(self):
        from ..context_menu import ContextMenu

        toolbar: Toolbar = omni.kit.widget.toolbar.get_instance()
        self.assertIsInstance(toolbar.context_menu, ContextMenu)

        toolbar.set_axis(ui.ToolBarAxis.X)
        toolbar.rebuild_toolbar()
        await ui_test.human_delay()

        toolbar.set_axis(ui.ToolBarAxis.Y)
        toolbar.rebuild_toolbar()
        await ui_test.human_delay()

        toolbar.subscribe_grab_mouse_pressed(None)
        
    async def test_toolbar_play_stop_button(self):
        toolbar: Toolbar = omni.kit.widget.toolbar.get_instance()
        timeline = omni.timeline.get_timeline_interface()

        # Click the stop button
        toolbar.get_widget("stop").call_clicked_fn()
        await ui_test.human_delay()
        self.assertTrue(timeline.is_stopped())
