# 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.input
import omni.appwindow
import omni.kit.app
import omni.kit.test
import omni.kit.window.toolbar
import omni.ui as ui
import omni.kit.ui_test as ui_test
from omni.ui.tests.test_base import OmniUiTest

from omni.kit.ui_test import emulate_mouse_move, emulate_mouse_move_and_click, Vec2
from omni.kit.ui_test.input import emulate_mouse
from carb.input import MouseEventType

from omni.kit.widget.toolbar.tests.helpers import reset_toolbar_settings
from omni.kit.widget.toolbar import WidgetGroup

GOLDEN_IMAGE_DIR = omni.kit.window.toolbar.get_data_path().joinpath("test", "golden_img").absolute()

class ToolbarDockingTest(OmniUiTest):
    async def setUp(self):
        await super().setUp()

        reset_toolbar_settings()

        self._main_dockspace = ui.Workspace.get_window("DockSpace")
        self._toolbar_handle = ui.Workspace.get_window(omni.kit.window.toolbar.Toolbar.WINDOW_NAME)
        self._toolbar_handle.undock()

        # Move mouse to origin
        await emulate_mouse_move(Vec2(0, 0))
        await ui_test.human_delay(10)

    async def test_docking_vertical(self):
        """Test vertical docking behavior."""
        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.LEFT)
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=GOLDEN_IMAGE_DIR, golden_img_name="v_dock.png")

    async def test_docking_horizontal(self):
        """Test horizontal docking behavior."""
        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.TOP)
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=GOLDEN_IMAGE_DIR, golden_img_name="h_dock.png")

    async def test_docking_vertical_context_menu(self):
        """Test vertical docking behavior showing context menu."""
        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.LEFT)
        await ui_test.human_delay(10)

        await emulate_mouse_move_and_click(Vec2(35, 135), right_click=True)
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=GOLDEN_IMAGE_DIR, golden_img_name="v_dock_menu.png")

    async def test_docking_horizontal_context_menu(self):
        """Test horizontal docking behavior showing context menu."""
        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.TOP)
        await ui_test.human_delay(10)

        await emulate_mouse_move_and_click(Vec2(125, 35), right_click=True)
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=GOLDEN_IMAGE_DIR, golden_img_name="h_dock_menu.png")

    async def test_docking_vertical_grab_context_menu(self):
        """Test vertical docking behavior showing grabber context menu."""
        self._toolbar_handle.dock_in(self._main_dockspace, ui.DockPosition.LEFT)
        await ui_test.human_delay(10)

        await emulate_mouse_move_and_click(Vec2(15, 35), right_click=True)
        await ui_test.human_delay(10)

        # sometimes the menu appears 1 pixel lower than expected
        await self.finalize_test(golden_img_dir=GOLDEN_IMAGE_DIR, golden_img_name="v_dock_grab_menu.png", threshold=0.07)

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

            def __init__(self):
                super().__init__()

            def get_style(self):
                return {}

            def create(self, default_size):
                button1 = ui.ToolButton(
                    name="test1",
                    width=default_size,
                    height=default_size,
                )

                button2 = ui.ToolButton(
                    name="test2",
                    width=default_size,
                    height=default_size,
                )

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

        widget_group = TestToolButtonGroup()

        toolbar = omni.kit.window.toolbar.get_instance()

        context = toolbar.acquire_toolbar_context('test')
        toolbar.add_widget(widget_group, 100)
        await ui_test.human_delay()
        self.assertIsNotNone(toolbar.get_widget('test1'))
        toolbar.remove_widget(widget_group)

        toolbar.add_custom_select_type('test', ['test'])
        toolbar.remove_custom_select('test')

        self.assertEqual(toolbar.get_context(), 'test')
        toolbar.release_toolbar_context(context)
