# Copyright (c) 2020, 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 carb.input
import carb.dictionary
import carb.settings
import omni.kit.context_menu
import omni.ui as ui
from carb.input import KeyboardInput as Key
from omni.kit.widget.toolbar.hotkey import Hotkey
from omni.kit.widget.toolbar.widget_group import WidgetGroup

from .models.transform_mode_model import LocalGlobalTransformModeModel, TransformModeModel

MOVE_TOOL_NAME = "Move"
ROTATE_TOOL_NAME = "Rotate"
SCALE_TOOL_NAME = "Scale"


class TransformButtonGroup(WidgetGroup):
    TRANSFORM_MOVE_MODE_SETTING = "/app/transform/moveMode"
    TRANSFORM_ROTATE_MODE_SETTING = "/app/transform/rotateMode"

    def __init__(self):
        super().__init__()
        self._input = carb.input.acquire_input_interface()
        self._settings = carb.settings.get_settings()
        self._settings.set_default_string(TransformModeModel.TRANSFORM_OP_SETTING, TransformModeModel.TRANSFORM_OP_MOVE)
        self._settings.set_default_string(
            self.TRANSFORM_MOVE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
        )
        self._settings.set_default_string(
            self.TRANSFORM_ROTATE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
        )

        self._move_op_model = LocalGlobalTransformModeModel(
            op=TransformModeModel.TRANSFORM_OP_MOVE, op_space_setting_path=self.TRANSFORM_MOVE_MODE_SETTING
        )
        self._rotate_op_model = LocalGlobalTransformModeModel(
            op=TransformModeModel.TRANSFORM_OP_ROTATE, op_space_setting_path=self.TRANSFORM_ROTATE_MODE_SETTING
        )
        self._scale_op_model = TransformModeModel(TransformModeModel.TRANSFORM_OP_SCALE)

        def on_hotkey_changed(hotkey: str, button, tool_name: str):
            button.tooltip = f"{tool_name} ({hotkey})"

        self._move_hotkey = Hotkey(
            "toolbar::move",
            Key.W,
            lambda: self._move_op_model.set_value(not self._move_op_model.get_value_as_bool()),
            lambda: self._move_button.enabled
            and self._is_in_context()
            and self._input.get_mouse_value(None, carb.input.MouseInput.RIGHT_BUTTON)
            == 0,  # when RMB is down, it's possible viewport WASD navigation is going on, and don't trigger it if W is pressed
            on_hotkey_changed_fn=lambda hotkey: on_hotkey_changed(hotkey, self._move_button, MOVE_TOOL_NAME),
        )
        self._rotate_hotkey = Hotkey(
            "toolbar::rotate",
            Key.E,
            lambda: self._rotate_op_model.set_value(not self._rotate_op_model.get_value_as_bool()),
            lambda: self._rotate_button.enabled and self._is_in_context()
            and self._input.get_mouse_value(None, carb.input.MouseInput.RIGHT_BUTTON)
            == 0,
            on_hotkey_changed_fn=lambda hotkey: on_hotkey_changed(hotkey, self._rotate_button, ROTATE_TOOL_NAME),
        )
        self._scale_hotkey = Hotkey(
            "toolbar::scale",
            Key.R,
            lambda: self._scale_op_model.set_value(True),
            lambda: self._scale_button.enabled and self._is_in_context(),
            on_hotkey_changed_fn=lambda hotkey: on_hotkey_changed(hotkey, self._scale_button, SCALE_TOOL_NAME),
        )

        self._register_context_menu()

    def get_style(self):
        style = {
            "Button.Image::move_op_global": {"image_url": "${glyphs}/toolbar_move_global.svg"},
            "Button.Image::move_op_local": {"image_url": "${glyphs}/toolbar_move_local.svg"},
            "Button.Image::rotate_op_global": {"image_url": "${glyphs}/toolbar_rotate_global.svg"},
            "Button.Image::rotate_op_local": {"image_url": "${glyphs}/toolbar_rotate_local.svg"},
            "Button.Image::scale_op": {"image_url": "${glyphs}/toolbar_scale.svg"},
        }
        return style

    def create(self, default_size):
        self._sub_move, self._move_button = self._create_local_global_button(
            self._move_op_model,
            f"{MOVE_TOOL_NAME} ({self._move_hotkey.get_as_string('W')})",
            "move_op",
            self.TRANSFORM_MOVE_MODE_SETTING,
            "move_op_global",
            "move_op_local",
            default_size,
        )
        self._sub_rotate, self._rotate_button = self._create_local_global_button(
            self._rotate_op_model,
            f"{ROTATE_TOOL_NAME} ({self._rotate_hotkey.get_as_string('E')})",
            "rotate_op",
            self.TRANSFORM_ROTATE_MODE_SETTING,
            "rotate_op_global",
            "rotate_op_local",
            default_size,
        )
        self._scale_button = ui.ToolButton(
            model=self._scale_op_model,
            name="scale_op",
            tooltip=f"{SCALE_TOOL_NAME} ({self._scale_hotkey.get_as_string('R')})",
            width=default_size,
            height=default_size,
            mouse_pressed_fn=lambda *_: self._acquire_toolbar_context(),
        )

        return {"move_op": self._move_button, "rotate_op": self._rotate_button, "scale_op": self._scale_button}

    def clean(self):
        super().clean()
        self._move_button = None
        self._rotate_button = None
        self._scale_button = None
        self._sub_move = None
        self._sub_rotate = None
        self._move_op_model.clean()
        self._move_op_model = None
        self._rotate_op_model.clean()
        self._rotate_op_model = None
        self._scale_op_model.clean()
        self._scale_op_model = None
        self._move_hotkey.clean()
        self._move_hotkey = None
        self._rotate_hotkey.clean()
        self._rotate_hotkey = None
        self._scale_hotkey.clean()
        self._scale_hotkey = None

        self._move_global_menu_entry = None
        self._move_local_menu_entry = None
        self._rotate_global_menu_entry = None
        self._rotate_local_menu_entry = None

    def _get_op_button_name(self, model, global_name, local_name):
        if model.get_op_space_mode() == LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL:
            return global_name
        else:  # pragma: no cover
            return local_name

    def _create_local_global_button(
        self, op_model, tooltip, button_name, op_setting_path, global_name, local_name, default_size
    ):

        op_button = ui.ToolButton(
            name=self._get_op_button_name(op_model, global_name, local_name),
            model=op_model,
            tooltip=tooltip,
            width=default_size,
            height=default_size,
            mouse_pressed_fn=lambda x, y, b, _: self._on_mouse_pressed(b, button_name),
            mouse_released_fn=lambda x, y, b, _: self._on_mouse_released(b),
            checked=op_model.get_value_as_bool(),
        )

        def on_op_button_value_change(model):
            op_button.name = self._get_op_button_name(model, global_name, local_name)

        return op_model.subscribe_value_changed_fn(on_op_button_value_change), op_button

    def _register_context_menu(self):
        context_menu = omni.kit.context_menu.get_instance()

        def is_button(objects: dict, button_name: str):  # pragma: no cover
            return objects.get("widget_name", None) == button_name

        def on_mode(objects: dict, op_setting_path: str, desired_value):  # pragma: no cover
            self._settings.set(op_setting_path, desired_value)

        def checked(objects: dict, op_setting_path: str, desired_value: str):  # pragma: no cover
            return self._settings.get(op_setting_path) == desired_value

        if context_menu:
            menu = {
                "name": "Global (W)",
                "show_fn": [lambda object: is_button(object, "move_op")],
                "onclick_fn": lambda object: on_mode(
                    object, self.TRANSFORM_MOVE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
                ),
                "checked_fn": lambda object: checked(
                    object, self.TRANSFORM_MOVE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
                ),
            }
            self._move_global_menu_entry = omni.kit.context_menu.add_menu(menu, "move_op", "omni.kit.widget.toolbar")

            menu = {
                "name": "Local (W)",
                "show_fn": [lambda object: is_button(object, "move_op")],
                "onclick_fn": lambda object: on_mode(
                    object, self.TRANSFORM_MOVE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_LOCAL
                ),
                "checked_fn": lambda object: checked(
                    object, self.TRANSFORM_MOVE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_LOCAL
                ),
            }
            self._move_local_menu_entry = omni.kit.context_menu.add_menu(menu, "move_op", "omni.kit.widget.toolbar")

            menu = {
                "name": "Global (E)",
                "show_fn": [lambda object: is_button(object, "rotate_op")],
                "onclick_fn": lambda object: on_mode(
                    object, self.TRANSFORM_ROTATE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
                ),
                "checked_fn": lambda object: checked(
                    object, self.TRANSFORM_ROTATE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_GLOBAL
                ),
            }
            self._rotate_global_menu_entry = omni.kit.context_menu.add_menu(
                menu, "rotate_op", "omni.kit.widget.toolbar"
            )

            menu = {
                "name": "Local (E)",
                "show_fn": [lambda object: is_button(object, "rotate_op")],
                "onclick_fn": lambda object: on_mode(
                    object, self.TRANSFORM_ROTATE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_LOCAL
                ),
                "checked_fn": lambda object: checked(
                    object, self.TRANSFORM_ROTATE_MODE_SETTING, LocalGlobalTransformModeModel.TRANSFORM_MODE_LOCAL
                ),
            }
            self._rotate_local_menu_entry = omni.kit.context_menu.add_menu(menu, "rotate_op", "omni.kit.widget.toolbar")
