# 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 omni.kit.context_menu
import omni.ui as ui
import pathlib
from pxr import Kind
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.select_mode_model import SelectModeModel
from .models.select_no_kinds_model import SelectNoKindsModel
from .models.select_include_ref_model import SelectIncludeRefModel
from .models.transform_mode_model import TransformModeModel

SELECT_PRIM_MODE_TOOL_NAME = "Prim"
SELECT_MODEL_MODE_TOOL_NAME = "Model"
SELECT_TOOL_NAME = "Select"
LIGHT_TYPES = "type:CylinderLight;type:DiskLight;type:DistantLight;type:DomeLight;type:GeometryLight;type:Light;type:RectLight;type:SphereLight"
EXTENSION_FOLDER_PATH = pathlib.Path(
    omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module("omni.kit.widget.toolbar")
)
SELECT_MODE_BUTTON_ENABLED_SETTING_PATH = "/exts/omni.kit.widget.toolbar/SelectionButton/SelectMode/enabled"


class SelectButtonGroup(WidgetGroup):
    def __init__(self):
        super().__init__()
        self._input = carb.input.acquire_input_interface()
        self._select_mode_model = SelectModeModel()
        self._select_no_kinds_model = SelectNoKindsModel()
        self._select_include_ref_model = SelectIncludeRefModel()
        self._select_op_model = TransformModeModel(TransformModeModel.TRANSFORM_OP_SELECT)
        self._select_op_menu_entry = None
        self._custom_types = []
        self._settings = carb.settings.get_settings()
        self._icon_path = str(EXTENSION_FOLDER_PATH / "data" / "icon")

        self._settings = carb.settings.get_settings()
        self._select_mode_button = None
        self._mode_hotkey = None
        if self._settings.get(SELECT_MODE_BUTTON_ENABLED_SETTING_PATH):
            def hotkey_change():
                current_mode = self._select_mode_model.get_value_as_string()
                if current_mode == "models" or current_mode == "kind:model.ALL":
                    self._select_mode_model.set_value("type:ALL")
                else:
                    self._select_mode_model.set_value("kind:model.ALL")
                self._update_selection_mode_button()

            def on_select_mode_hotkey_changed(hotkey: str):
                self._select_mode_button.tooltip = self._get_select_mode_tooltip()

            self._mode_hotkey = Hotkey(
                "toolbar::select_mode",
                Key.T,
                hotkey_change,
                lambda: self._select_mode_button.enabled and self._is_in_context(),
                on_hotkey_changed_fn=on_select_mode_hotkey_changed,
            )

        # only replace the tooltip part so that "Paint Select" registered from select brush can still have correct label
        def on_select_hotkey_changed(hotkey: str):
            self._select_op_button.tooltip = (
                self._select_op_button.tooltip.rsplit("(", 1)[0] + f"({self._select_hotkey.get_as_string('Q')})"
            )

        self._select_hotkey = Hotkey(
            "toolbar::select",
            Key.Q,
            lambda: self._select_op_button.model.set_value(True),
            lambda: self._select_op_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_select_hotkey_changed(hotkey),
        )
        self._sub_menu = []
        self._register_context_menu()

    def get_style(self):
        style = {
            "Button.Image::select_mode": {"image_url": "${glyphs}/toolbar_select_prim.svg"},
            "Button.Image::select_op_models": {"image_url": "${glyphs}/toolbar_select_models.svg"},
            "Button.Image::select_op_prims": {"image_url": "${glyphs}/toolbar_select_prims.svg"},

            "Button.Image::all_prim_types": {"image_url": f"{self._icon_path}/all_prim_types.svg"},
            "Button.Image::meshes": {"image_url": f"{self._icon_path}/meshes.svg"},
            "Button.Image::lights": {"image_url": f"{self._icon_path}/lights.svg"},
            "Button.Image::cameras": {"image_url": f"{self._icon_path}/cameras.svg"},
            "Button.Image::model_kinds": {"image_url": f"{self._icon_path}/model_kinds.svg"},
            "Button.Image::assembly": {"image_url": f"{self._icon_path}/assembly.svg"},
            "Button.Image::group": {"image_url": f"{self._icon_path}/group.svg"},
            "Button.Image::component": {"image_url": f"{self._icon_path}/component.svg"},
            "Button.Image::sub_component": {"image_url": f"{self._icon_path}/sub_component.svg"},
            "Button.Image::payload_reference": {"image_url": f"{self._icon_path}/payload_reference.svg"},
            "Button.Image::custom_kind": {"image_url": f"{self._icon_path}/custom_kind.svg"},
            "Button.Image::custom_type": {"image_url": f"{self._icon_path}/custom_type.svg"},
        }
        return style

    def create(self, default_size):

        def select_enabled():  # pragma: no cover (Unused, it's commened out below - I am just adding coverage! '^^)
            return True

        def on_select_mode_change(model):
            self._update_selection_mode_button()

        result = {}
        if self._settings.get(SELECT_MODE_BUTTON_ENABLED_SETTING_PATH):
            self._select_mode_button = ui.ToolButton(
                model=self._select_mode_model,
                name="select_mode",
                tooltip=self._get_select_mode_tooltip(),
                width=default_size,
                height=default_size,
                # checked=select_enabled,
                mouse_pressed_fn=lambda x, y, b, _: self._on_mouse_pressed(b, "select_mode"),
                mouse_released_fn=lambda x, y, b, _: self._on_mouse_released(b),
            )
            self._model_changed_sub = self._select_mode_model.subscribe_value_changed_fn(on_select_mode_change)
            result["select_mode"] = self._select_mode_button
        with ui.ZStack(width=0, height=0):
            self._select_op_button = ui.ToolButton(
                model=self._select_op_model,
                name=self._get_select_op_button_name(self._select_mode_model),
                tooltip=self._get_select_tooltip(),
                width=default_size,
                height=default_size,
                checked=self._select_op_model.get_value_as_bool(),
                mouse_pressed_fn=lambda x, y, b, _: self._on_mouse_pressed(b, "select_op"),
                mouse_released_fn=lambda x, y, b, _: self._on_mouse_released(b),
            )
            # Follow the pattern here (ZStack + _build_flyout_indicator) for other buttons if they want dynamic flyout indicator
            self._select_op_menu_sub = self._build_flyout_indicator(default_size, default_size, "select_op")
            result["select_op"] = self._select_op_button

        self._update_selection_mode_button()
        return result

    def clean(self):
        super().clean()
        self._select_mode_button = None
        self._select_op_button = None
        self._select_mode_model.clean()
        self._select_mode_model = None
        self._select_no_kinds_model.clean()
        self._select_no_kinds_model = None
        self._select_include_ref_model.clean()
        self._select_include_ref_model = None
        self._select_op_model.clean()
        self._select_op_model = None
        self._model_changed_sub = None
        if self._mode_hotkey:
            self._mode_hotkey.clean()
        self._mode_hotkey = None
        self._select_hotkey.clean()
        self._select_hotkey = None
        self._select_mode_menu_entry.release()
        self._select_mode_menu_entry = None
        if self._select_op_menu_entry:
            self._select_op_menu_entry.release()
        self._select_op_menu_entry = None
        self._select_op_menu_sub = None
        for menu in self._sub_menu:
            menu.release()
        self._sub_menu = []

    def _get_select_op_button_name(self, model):
        return "select_op_prims" if model.get_value_as_bool() else "select_op_models"

    def _get_select_tooltip(self):
        return f"{SELECT_TOOL_NAME} ({self._select_hotkey.get_as_string('Q')})"

    def _get_select_mode_button_name(self):
        mode_name = {
            "type:ALL": "all_prim_types",
            "type:Mesh": "meshes",
            LIGHT_TYPES: "lights",
            "type:Camera": "cameras",
            "kind:model.ALL": "model_kinds",
            "kind:assembly": "assembly",
            "kind:group": "group",
            "kind:component": "component",
            "kind:subcomponent": "sub_component",
            "ref:reference;ref:payload": "custom_type",
        }
        mode = self._select_mode_model.get_value_as_string()
        if mode in mode_name:
            return mode_name[mode]
        else:  # pragma: no cover
            if mode[:4] == "kind":
                return "custom_kind"
            else:
                return "custom_type"

    def _get_select_mode_tooltip(self):
        mode_data = {
            "type:ALL": "All Prim Types",
            "type:Mesh": "Meshes",
            LIGHT_TYPES: "Lights",
            "type:Camera": "Camera",
            "kind:model.ALL": "All Model Kinds",
            "kind:assembly": "Assembly",
            "kind:group": "Group",
            "kind:component": "Component",
            "kind:subcomponent": "Subcomponent",
            "ref:reference;ref:payload": "Subcomponent",
        }
        mode = self._select_mode_model.get_value_as_string()
        if mode in mode_data:
            tooltip = mode_data[mode]
        else:  # pragma: no cover
            if mode[:4] == "kind":
                tooltip = "Custom Kind"
            else:
                tooltip = "Custom type"

        return tooltip + f" ({self._mode_hotkey.get_as_string('T')})"

    def _usd_kinds(self):
        return ["model", "assembly", "group", "component", "subcomponent"]

    def _enable_no_kinds_option(self, b):  # pragma: no cover
        mode = self._select_mode_model.get_value_as_string()
        return mode[:4] == "kind"

    def _usd_kinds_display(self):  # pragma: no cover
        return self._usd_kinds()[1:]

    def _plugin_kinds(self):
        all_kinds = set(Kind.Registry.GetAllKinds())
        return all_kinds - set(self._usd_kinds())

    def _create_selection_list(self, selections):
        results = ""
        delimiter = ";"
        first_selection = True
        for s in selections:
            if not first_selection:  # pragma: no cover (This seemingly can never happen, since first_selection is initialized with True oO)
                first_selection = True
            else:
                results += delimiter
            results += s
        return results

    def _update_selection_mode_button(self):
        if self._select_mode_button:
            self._select_mode_button.name = self._get_select_mode_button_name()
            self._select_mode_button.tooltip = self._get_select_mode_tooltip()

        if self._select_op_button:
            self._select_op_button.name = self._get_select_op_button_name(self._select_mode_model)

    def _update_select_menu(self):
        for menu in self._sub_menu:
            menu.release()

        self._sub_menu = []

        def create_menu_entry(name, value, add_hotkey=False):
            menu = {
                "name": name,
                "onclick_fn": lambda object: self._select_mode_model.set_value(value),
                "checked_fn": lambda object, name=name, value=value: self._select_mode_model.get_value_as_string() == value,
                "additional_kwargs": {
                    "style": {"Menu.Item.CheckMark": {"image_url": omni.kit.context_menu.style.get_radio_mark_url()}}
                },
            }
            if add_hotkey:
                menu["hotkey"] = (carb.input.KeyboardInput.T)

            return omni.kit.context_menu.add_menu(menu, "select_mode", "omni.kit.widget.toolbar")

        def create_header(header=""):
            menu = {
                "header": header,
                "name": "",
            }
            return omni.kit.context_menu.add_menu(menu, "select_mode", "omni.kit.widget.toolbar")

        self._sub_menu.append(create_header(header="Select by Type"))
        self._sub_menu.append(create_menu_entry("All Prim Types", "type:ALL", True))
        self._sub_menu.append(create_menu_entry("Meshes", "type:Mesh"))
        self._sub_menu.append(create_menu_entry("Lights", LIGHT_TYPES))
        self._sub_menu.append(create_menu_entry("Camera", "type:Camera"))

        if self._custom_types:
            for name, types in self._custom_types:
                self._sub_menu.append(create_menu_entry(name, types))

        self._sub_menu.append(create_header(header="Select by Model Kind"))
        self._sub_menu.append(create_menu_entry("All Model Kinds", "kind:model.ALL"))

        self._sub_menu.append(create_menu_entry("Assembly", "kind:assembly"))
        self._sub_menu.append(create_menu_entry("Group", "kind:group"))
        self._sub_menu.append(create_menu_entry("Component", "kind:component"))
        self._sub_menu.append(create_menu_entry("Subcomponent", "kind:subcomponent"))

        plugin_kinds = self._plugin_kinds()
        if (len(plugin_kinds)) > 0:
            for k in plugin_kinds:  # pragma: no cover
                self._sub_menu.append(create_menu_entry(str(k), str(key).capitalize(), f"kind:{k}"))

        self._sub_menu.append(create_header(header=""))

        menu = {
            "name": "Include Prims with no Kind",
            "onclick_fn": lambda object: self._select_no_kinds_model.set_value(not self._select_no_kinds_model.get_value_as_bool()),
            "checked_fn": lambda object: self._select_no_kinds_model.get_value_as_bool(),
            "enabled_fn": self._enable_no_kinds_option,
        }
        self._sub_menu.append(omni.kit.context_menu.add_menu(menu, "select_mode", "omni.kit.widget.toolbar"))

        menu = {
            "name": "Include References and Payloads",
            "onclick_fn": lambda object: self._select_include_ref_model.set_value(not self._select_include_ref_model.get_value_as_bool()),
            "checked_fn": lambda object: self._select_include_ref_model.get_value_as_bool(),
            "enabled_fn": self._enable_no_kinds_option,
        }
        self._sub_menu.append(omni.kit.context_menu.add_menu(menu, "select_mode", "omni.kit.widget.toolbar"))

    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

        if context_menu:
            def on_clicked_select():  # pragma: no cover
                self._select_op_model.set_value(True)
                self._select_op_button.name = self._get_select_op_button_name(self._select_mode_model)
                self._select_op_button.set_tooltip(self._get_select_tooltip())
                self._select_op_button.model = self._select_op_model

            menu = {
                "name": "Select",
                "show_fn": [lambda object: is_button(object, "select_op")],
                "onclick_fn": lambda object: on_clicked_select(),
                "checked_fn": lambda object: self._select_op_button.name == "select_op_prims"
                or self._select_op_button.name == "select_op_models",
            }
            self._select_mode_menu_entry = omni.kit.context_menu.add_menu(menu, "select_op", "omni.kit.widget.toolbar")

            if self._settings.get(SELECT_MODE_BUTTON_ENABLED_SETTING_PATH):
                menu = {
                    "name": "Select Mode",
                    "show_fn": [lambda object: is_button(object, "select_mode")],
                }
                self._select_op_menu_entry = omni.kit.context_menu.add_menu(menu, "select_mode", "omni.kit.widget.toolbar")

            self._update_select_menu()

    def add_custom_select_type(self, entry_name: str, selection_types: list):
        selection = []
        for s in selection_types:
            selection.append(f"type:{s}")
        selection_string = self._create_selection_list(selection)
        self._custom_types.append((entry_name, selection_string))
        self._update_select_menu()

    def remove_custom_select(self, entry_name):
        for index, (entry, selection_string) in enumerate(self._custom_types):
            if entry == entry_name:
                del self._custom_types[index]
                self._update_select_menu()
                return
