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

__all__ = ["DisplayMenuContainer"]

from omni.kit.viewport.menubar.core import (
    IconMenuDelegate,
    SettingModel,
    ViewportMenuContainer,
    CategoryMenuContainer,
    SelectableMenuItem,
    SimpleCategoryModel,
    CategoryStateItem,
    BaseCategoryItem,
    CategoryCustomItem,
    CategoryCollectionItem
)
from .style import UI_STYLE

import carb
import carb.settings
import omni.ui as ui
import omni.kit.app
import omni.usd

from functools import partial
from typing import Dict, List, Optional

SHOW_BY_TYPE_EXCLUDE_LIST = "/exts/omni.kit.viewport.menubar.display/showByType/exclude_list"
HEADS_UP_CATEGORY_NAME = "Heads Up Display"
SHOW_BY_TYPE_CATEGORY_NAME = "Show By Type"
SHOW_BY_PURPOSE_CATEGORY_NAME = "Show By Purpose"
DEFAULT_CATEGORIES = [HEADS_UP_CATEGORY_NAME, SHOW_BY_TYPE_CATEGORY_NAME, SHOW_BY_PURPOSE_CATEGORY_NAME]
DEFAULT_SECTION = "default"


def _make_viewport_setting(viewport_api_id: str, setting: str):
    return f"/persistent/app/viewport/{viewport_api_id}/{setting}/visible"


class DisplayMenuContainer(ViewportMenuContainer):
    """The menu with the visibility settings"""

    def __init__(self):
        super().__init__(
            name="Display",
            delegate=IconMenuDelegate("Display"),
            visible_setting_path="/exts/omni.kit.viewport.menubar.display/visible",
            order_setting_path="/exts/omni.kit.viewport.menubar.display/order",
            style=UI_STYLE
        )

        self._root_menu: Optional[ui.Menu] = None
        self._category_models: Dict[str, SimpleCategoryModel] = {}
        self._custom_settings: List[List[str, str]] = []
        self._custom_category_items: Dict[str, List[BaseCategoryItem]] = {}
        self._section_categories: Dict[str, List[str]] = {}
        self._section_categories[DEFAULT_SECTION] = DEFAULT_CATEGORIES[:]  # Copy the default categories list

    def destroy(self):
        super().destroy()

    def register_custom_setting(self, text: str, setting_path: str):
        self._custom_settings.append((text, setting_path))
        if self._root_menu:
            self._root_menu.invalidate()

    def deregister_custom_setting(self, text: str):
        found = [item for item in self._custom_settings if item[0] == text]
        if found:
            for item in found:
                self._custom_settings.remove(item)
            if self._root_menu:
                self._root_menu.invalidate()

    def register_custom_category_item(self, category: str, item: BaseCategoryItem, section: str):
        is_top_category = False
        if category not in DEFAULT_CATEGORIES and category not in self._category_models:
            if item.text == category and isinstance(item, CategoryCollectionItem):
                self._category_models[category] = SimpleCategoryModel(category, root=item)
                is_top_category = True
            else:
                self._category_models[category] = SimpleCategoryModel(category)

        if category not in self._custom_category_items:
            self._custom_category_items[category] = []
        if section not in self._section_categories:
            self._section_categories[section] = []

        if not is_top_category:
            self._custom_category_items[category].append(item)

        if category not in self._section_categories[section]:
            self._section_categories[section].append(category)
        if self._root_menu:
            self._root_menu.invalidate()

    def deregister_custom_category_item(self, category: str, item: BaseCategoryItem):
        if category in self._custom_category_items:
            if item in self._custom_category_items[category]:
                self._custom_category_items[category].remove(item)

        if category not in DEFAULT_CATEGORIES:
            if (item.text == category and isinstance(item, CategoryCollectionItem)) or len(self._custom_category_items[category]) == 0:
                del self._category_models[category]
                # Now clean up section
                sections = list(self._section_categories.keys())
                for section in sections:
                    if category in self._section_categories[section]:
                        self._section_categories[section].remove(category)
                        if len(self._section_categories[section]) == 0:
                            del self._section_categories[section]

        if self._root_menu:
            self._root_menu.invalidate()

    def build_fn(self, viewport_context: dict):
        self._root_menu = ui.Menu(self.name, delegate=self._delegate,
                                  on_build_fn=partial(self._build_menu_items, viewport_context),
                                  style=self._style)

    def _build_menu_items(self, viewport_context: dict, *args, **kwargs):
        viewport = viewport_context.get("viewport_api")
        viewport_api_id: str = str(viewport.id)
        settings = carb.settings.get_settings()

        show_by_type_items: list[BaseCategoryItem] = [
            CategoryStateItem("Cameras", setting_path=_make_viewport_setting(viewport_api_id, "scene/cameras")),
            CategoryStateItem("Lights", setting_path=_make_viewport_setting(viewport_api_id, "scene/lights")),
            CategoryStateItem("Skeletons", setting_path=_make_viewport_setting(viewport_api_id, "scene/skeletons")),
            CategoryStateItem("Audio", setting_path=_make_viewport_setting(viewport_api_id, "scene/audio")),
        ]
        if (exclude_list := settings.get(SHOW_BY_TYPE_EXCLUDE_LIST)):
            show_by_type_items = [item for item in show_by_type_items if item.text not in exclude_list]

        # 105.1: Support alternate label of memory (i.e. "Host Memory", "Process Memory", "Memory")
        # Defaults to pre 105.1 label (Host Memory) when not specified
        mem_label = settings.get("/exts/omni.kit.viewport.window/hud/hostMemory/label")
        if mem_label is None:
            mem_label = "Host"

        default_category_models = {
            HEADS_UP_CATEGORY_NAME: SimpleCategoryModel(
                HEADS_UP_CATEGORY_NAME,
                [
                    CategoryStateItem("FPS", setting_path=_make_viewport_setting(viewport_api_id, "hud/renderFPS")),
                    CategoryStateItem("Device Memory", setting_path=_make_viewport_setting(viewport_api_id, "hud/deviceMemory")),
                    CategoryStateItem(f"{mem_label} Memory", setting_path=_make_viewport_setting(viewport_api_id, "hud/hostMemory")),
                    CategoryStateItem("Resolution", setting_path=_make_viewport_setting(viewport_api_id, "hud/renderResolution")),
                    CategoryStateItem("Progress", setting_path=_make_viewport_setting(viewport_api_id, "hud/renderProgress")),
                ]
            ),

            SHOW_BY_TYPE_CATEGORY_NAME: SimpleCategoryModel(
                SHOW_BY_TYPE_CATEGORY_NAME,
                show_by_type_items
            ),

            SHOW_BY_PURPOSE_CATEGORY_NAME: SimpleCategoryModel(
                SHOW_BY_PURPOSE_CATEGORY_NAME,
                [
                    CategoryStateItem("Guide", setting_path="/persistent/app/hydra/displayPurpose/guide"),
                    CategoryStateItem("Proxy", setting_path="/persistent/app/hydra/displayPurpose/proxy"),
                    CategoryStateItem("Render", setting_path="/persistent/app/hydra/displayPurpose/render"),
                ]
            )
        }
        self._category_models.update(default_category_models)

        # XXX: These add_item calls currently must occur to add the separator!
        self._category_models[SHOW_BY_TYPE_CATEGORY_NAME].add_item(CategoryCustomItem(
            "Meshes",
            lambda: SelectableMenuItem("Meshes",
                                       SettingModel(setting_path=_make_viewport_setting(viewport_api_id, "scene/meshes")))
        ))
        self._category_models[HEADS_UP_CATEGORY_NAME].add_item(CategoryCustomItem(
            "Camera Speed",
            lambda: SelectableMenuItem("Camera Speed",
                                       SettingModel(_make_viewport_setting(viewport_api_id, "hud/cameraSpeed")))
        ))

        identifier = "omni.kit.viewport.menubar.display"

        # Create default section categories first
        for name in self._section_categories[DEFAULT_SECTION]:
            model = self._category_models[name]
            if name in self._custom_category_items:
                for item in self._custom_category_items[name]:
                    model.add_item(item)

            # XXX: Workaround nested creation of these items not being able to trigger an action!
            trigger_fns = None
            if name == SHOW_BY_TYPE_CATEGORY_NAME:
                icon_click_id = f"{identifier}.{name}.{name}"  # Left-most check/mixed icon was toggled
                trigger_fns = {
                    "Cameras": partial(self.__trigger_action, "toggle_camera_visibility", viewport_api=viewport),
                    "Lights": partial(self.__trigger_action, "toggle_light_visibility", viewport_api=viewport),
                    "Skeletons": partial(self.__trigger_action, "toggle_skeleton_visibility", viewport_api=viewport),
                    "Audio": partial(self.__trigger_action, "toggle_audio_visibility", viewport_api=viewport),
                    "Meshes": partial(self.__trigger_action, "toggle_mesh_visibility", viewport_api=viewport),
                    icon_click_id: partial(self.__trigger_action, "toggle_show_by_type_visibility", viewport_api=viewport),
                }

            CategoryMenuContainer(model, identifier=f"{identifier}.{name}", trigger_fns=trigger_fns)

        # Now iterate named sections, with a separator for each.
        for section, categories in self._section_categories.items():
            if section is DEFAULT_SECTION:
                continue
            ui.Separator(text=section)
            for name in categories:
                model = self._category_models[name]
                if name in self._custom_category_items:
                    for item in self._custom_category_items[name]:
                        model.add_item(item)
                CategoryMenuContainer(model, identifier=f"{identifier}.{name}")

        ui.Separator()

        # This currently is just easier tied to legacy global setting
        SelectableMenuItem("Selection Outline", SettingModel(_make_viewport_setting(viewport_api_id, "guide/selection")),
            triggered_fn=partial(self.__trigger_action, "toggle_selection_hilight_visibility", viewport_api=viewport),
            trigger_will_set_model=True
        )
        SelectableMenuItem("Axis", SettingModel(_make_viewport_setting(viewport_api_id, "guide/axis")),
            triggered_fn=partial(self.__trigger_action, "toggle_axis_visibility", viewport_api=viewport),
            trigger_will_set_model=True
        )
        SelectableMenuItem("Grid", SettingModel(_make_viewport_setting(viewport_api_id, "guide/grid")),
            triggered_fn=partial(self.__trigger_action, "toggle_grid_visibility", viewport_api=viewport),
            trigger_will_set_model=True
        )

        # Custom display settings
        if self._custom_settings:
            ui.Separator()
            for (text, setting_path) in self._custom_settings:
                SelectableMenuItem(text, SettingModel(setting_path))

    def __trigger_action(self, action: str, *args, **kwargs):
        import omni.kit.actions.core
        action_registry = omni.kit.actions.core.get_action_registry()
        if action_registry:
            exc_action = action_registry.get_action("omni.kit.viewport.actions", action)
            if exc_action:
                exc_action.execute(*args, **kwargs)
            else:
                carb.log_error(f"Could not find action to run: '{action}'")
        else:
            carb.log_error(f"Could not get action_registry to run '{action}")
