# Copyright (c) 2022, 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.
#

from pathlib import Path
from typing import List
import traceback

import carb
import carb.dictionary
import carb.events
import carb.settings
import omni.kit.app
import omni.kit.context_menu
import omni.usd
from omni.kit.manipulator.tool.snap import SnapToolButton
from omni.kit.manipulator.transform.manipulator import TransformManipulator
from omni.kit.manipulator.transform.toolbar_tool import SimpleToolButton
from omni.kit.manipulator.transform.types import Operation

from .settings_constants import Constants as prim_c
from .tool_models import LocalGlobalModeModel
from .toolbar_registry import get_toolbar_registry

ICON_FOLDER_PATH = Path(
    f"{omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)}/data/icons"
)

TOOLS_ENABLED_SETTING_PATH = "/exts/omni.kit.manipulator.prim/tools/enabled"


class SelectionPivotTool(SimpleToolButton):

    # menu entry only needs to register once
    __menu_entries = []

    @classmethod
    def register_menu(cls):
        def build_placement_setting_entry(setting: str):
            menu = {
                "name": setting,
                "checked_fn": lambda _: carb.settings.get_settings().get(prim_c.MANIPULATOR_PLACEMENT_SETTING)
                == setting,
                "onclick_fn": lambda _: carb.settings.get_settings().set(prim_c.MANIPULATOR_PLACEMENT_SETTING, setting),
            }
            return menu

        menu = [
            build_placement_setting_entry(s)
            for s in [
                prim_c.MANIPULATOR_PLACEMENT_LAST_PRIM_PIVOT,
                prim_c.MANIPULATOR_PLACEMENT_BBOX_BASE,
                prim_c.MANIPULATOR_PLACEMENT_BBOX_CENTER,
                prim_c.MANIPULATOR_PLACEMENT_SELECTION_CENTER,
                prim_c.MANIPULATOR_PLACEMENT_PICK_REF_PRIM,
            ]
        ]

        for item in menu:
            cls.__menu_entries.append(omni.kit.context_menu.add_menu(item, "sel_pivot", "omni.kit.manipulator.prim"))

    @classmethod
    def unregister_menu(cls):
        for sub in cls.__menu_entries:
            sub.release()
        cls.__menu_entries.clear()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # TODO assume the context is "" for VP1
        self._usd_context = omni.usd.get_context(self._toolbar_payload.get("usd_context_name", ""))
        self._selection = self._usd_context.get_selection()

        enabled_img_url = f"{ICON_FOLDER_PATH}/pivot_location.svg"

        self._build_widget(
            button_name="sel_pivot",
            enabled_img_url=enabled_img_url,
            model=None,
            menu_index="sel_pivot",
            menu_extension_id="omni.kit.manipulator.prim",
            no_toggle=True,
            menu_on_left_click=True,
            tooltip="Selection Pivot Placement",
        )

        # order=1 to register after prim manipulator so _on_selection_changed gets updated len of xformable_prim_paths
        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop_by_type(
            int(omni.usd.StageEventType.SELECTION_CHANGED),
            self._on_stage_selection_event,
            name="SelectionPivotTool stage event",
            order=1,
        )

        if self._usd_context.get_stage_state() == omni.usd.StageState.OPENED:
            self._on_selection_changed()

    def destroy(self):
        super().destroy()
        self._stage_event_sub = None
        if self._model:
            self._model.destroy()
            self._model = None

    @classmethod
    def can_build(cls, manipulator: TransformManipulator, operation: Operation) -> bool:
        return operation == Operation.TRANSLATE or operation == Operation.ROTATE or operation == Operation.SCALE

    def _on_stage_selection_event(self, event: carb.events.IEvent):
        self._on_selection_changed()

    def _on_selection_changed(self):
        selected_xformable_paths_count = (
            len(self._manipulator.model.xformable_prim_paths) if self._manipulator.model else 0
        )

        if self._stack:
            visible = selected_xformable_paths_count > 0
            if self._stack.visible != visible:
                self._stack.visible = visible
                self._manipulator.refresh_toolbar()


class LocalGlobalTool(SimpleToolButton):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        if self._operation == Operation.TRANSLATE:
            setting_path = "/app/transform/moveMode"
        elif self._operation == Operation.ROTATE:
            setting_path = "/app/transform/rotateMode"
        else:
            raise RuntimeError("Invalid operation")

        self._model = LocalGlobalModeModel(setting_path)
        enabled_img_url = f"{ICON_FOLDER_PATH}/transformspace_global_dark.svg"
        disabled_img_url = f"{ICON_FOLDER_PATH}/transformspace_local_dark.svg"

        self._build_widget(
            button_name="local_global",
            enabled_img_url=enabled_img_url,
            disabled_img_url=disabled_img_url,
            model=self._model,
            tooltip="Current Transform Space: World",
            disabled_tooltip="Current Transform Space: Local"
        )

    def destroy(self):
        super().destroy()
        if self._model:
            self._model.destroy()
            self._model = None

    @classmethod
    def can_build(cls, manipulator: TransformManipulator, operation: Operation) -> bool:
        return operation == Operation.TRANSLATE or operation == Operation.ROTATE


class PrimManipTools:
    def __init__(self):
        self._settings = carb.settings.get_settings()
        self._dict = carb.dictionary.get_dictionary()
        self._toolbar_reg = get_toolbar_registry()

        SelectionPivotTool.register_menu()  # one time startup
        self._builtin_tool_classes = {
            "prim:1space": LocalGlobalTool,
            "prim:2snap": SnapToolButton,
            "prim:3sel_pivot": SelectionPivotTool,
        }
        self._registered_tool_ids: List[str] = []

        self._sub = self._settings.subscribe_to_node_change_events(TOOLS_ENABLED_SETTING_PATH, self._on_setting_changed)
        if self._settings.get(TOOLS_ENABLED_SETTING_PATH) is True:
            self._register_tools()

        self._pivot_button_group = None

        # register pivot placement to "factory explorer" toolbar
        app_name = omni.kit.app.get_app().get_app_filename()
        if "omni.factory_explorer" in app_name:
            # add manu entry to factory explorer toolbar
            manager = omni.kit.app.get_app().get_extension_manager()
            self._hooks = manager.subscribe_to_extension_enable(
                lambda _: self._register_main_toolbar_button(),
                lambda _: self._unregister_main_toolbar_button(),
                ext_name="omni.explore.toolbar",
                hook_name="omni.kit.manipulator.prim.pivot_placement listener",
            )

    def __del__(self):
        self.destroy()

    def destroy(self):
        self._hooks = None
        if self._pivot_button_group is not None:
            self._unregister_main_toolbar_button()

        self._unregister_tools()
        if self._sub is not None:
            self._settings.unsubscribe_to_change_events(self._sub)
            self._sub = None

        SelectionPivotTool.unregister_menu()  # one time shutdown

    def _register_tools(self):
        for id, tool_class in self._builtin_tool_classes.items():
            self._toolbar_reg.register_tool(tool_class, id)
            self._registered_tool_ids.append(id)

    def _unregister_tools(self):
        for id in self._registered_tool_ids:
            self._toolbar_reg.unregister_tool(id)

        self._registered_tool_ids.clear()

    def _on_setting_changed(self, item, event_type):
        enabled = self._dict.get(item)
        if enabled and not self._registered_tool_ids:
            self._register_tools()
        elif not enabled:
            self._unregister_tools()

    def _register_main_toolbar_button(self):  # pragma: no cover
        # Can't be tested from within kit since omni.explore.toolbar is outside of kit repo
        try:
            if not self._pivot_button_group:
                import omni.explore.toolbar

                from .pivot_button_group import PivotButtonGroup

                self._pivot_button_group = PivotButtonGroup()

                explorer_toolbar_ext = omni.explore.toolbar.get_toolbar_instance()
                explorer_toolbar_ext.toolbar.add_widget_group(self._pivot_button_group, 1)

                # also add to MODIFY_TOOLS list, when menu Modify is selected
                # all items in toolbar will be cleared and re-added from MODIFY_TOOLS list
                explorer_tools = omni.explore.toolbar.groups.MODIFY_TOOLS
                explorer_tools.append(self._pivot_button_group)
        except Exception:
            carb.log_warn(traceback.format_exc())

    def _unregister_main_toolbar_button(self):  # pragma: no cover
        # Can't be tested from within kit since omni.explore.toolbar is outside of kit repo
        try:
            if self._pivot_button_group:
                import omni.explore.toolbar

                explorer_toolbar_ext = omni.explore.toolbar.get_toolbar_instance()
                explorer_toolbar_ext.toolbar.remove_widget_group(self._pivot_button_group)

                # remove from MODIFY_TOOLS list as well
                explorer_tools = omni.explore.toolbar.groups.MODIFY_TOOLS
                if self._pivot_button_group in explorer_tools:
                    explorer_tools.remove(self._pivot_button_group)

                self._pivot_button_group.clean()
                self._pivot_button_group = None

        except Exception:
            carb.log_warn(traceback.format_exc())
