# 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__ = ['ViewportExtension']

import omni.ext
import omni.kit.menu.utils
from omni.ui import Workspace
from omni.kit.menu.utils import MenuItemDescription
from .window import ViewportWindow
from .dragdrop.usd_file_drop_delegate import UsdFileDropDelegate
from .dragdrop.usd_prim_drop_delegate import UsdShadeDropDelegate
from .dragdrop.material_file_drop_delegate import MaterialFileDropDelegate
from .dragdrop.audio_file_drop_delegate import AudioFileDropDelegate
from .viewport_actions import register_actions, deregister_actions
import carb

from typing import Dict, List, Optional

DEFAULT_VIEWPORT_NAME = '/exts/omni.kit.viewport.window/startup/windowName'
DEFAULT_VIEWPORT_NO_OPEN = '/exts/omni.kit.viewport.window/startup/disableWindowOnLoad'
DEFAULT_VIEWPORT_HIDE_TAB = '/exts/omni.kit.viewport.window/startup/dockTabInvisible'


class _MenuEntry:
    def __init__(self, window_name: str, open_window: bool = False, index: int = 0):
        menu_label = f'{window_name} {index + 1}' # if index else window_name
        self.__window_name = menu_label if index else window_name
        self.__window = None
        self.__menu_entry = None
        self.__action_name = None

        Workspace.set_show_window_fn(self.__window_name, lambda b: self.__show_window(None, b))
        if open_window:
            Workspace.show_window(self.__window_name)

        # add actions
        self.__action_name = register_actions("omni.kit.viewport.window", index, self.__show_window, self.__is_window_visible)

        # add menu
        if self.__action_name:
            self.__menu_entry = [
                MenuItemDescription(name="Viewport", sub_menu=
                    [MenuItemDescription(
                        name=f"{menu_label}",
                        ticked=True,
                        ticked_fn=self.__is_window_visible,
                        onclick_action=("omni.kit.viewport.window", self.__action_name),
                    )]
                )
            ]
            omni.kit.menu.utils.add_menu_items(self.__menu_entry, name="Window")

    def __del__(self):
        self.destroy()

    def destroy(self):
        if self.__window:
            Workspace.set_show_window_fn(self.__window_name, None)
            self.__window.set_visibility_changed_fn(None)
            self.__window.destroy()
            self.__window = None

        # remove menu
        if self.__menu_entry:
            omni.kit.menu.utils.remove_menu_items(self.__menu_entry, name="Window")

        # remove actions
        if self.__action_name:
            deregister_actions("omni.kit.viewport.window", self.__action_name)

    def __dock_window(self):
        if not self.__window:
            return
        settings = carb.settings.get_settings()
        if bool(settings.get("/app/docks/disabled")):
            return

        async def wait_for_window(window_name: str, dock_name: str, position: omni.ui.DockPosition, ratio: float = 1):
            dockspace = Workspace.get_window(dock_name)
            window = Workspace.get_window(window_name)
            if (window is None) or (dockspace is None):
                frames = 3
                while ((window is None) or (dockspace is None)) and frames:
                    await omni.kit.app.get_app().next_update_async()
                    dockspace = Workspace.get_window(dock_name)
                    window = Workspace.get_window(window_name)
                    frames = frames - 1

            if window and dockspace:
                # When docking to existing Viewport, preserve tab-visibility for both, otherwise defer to setting
                if position == omni.ui.DockPosition.SAME:
                    dock_tab_bar_visible = not bool(settings.get(DEFAULT_VIEWPORT_HIDE_TAB))
                else:
                    dock_tab_bar_visible = dockspace.dock_tab_bar_visible

                window.deferred_dock_in(dock_name)
                # Split the tab to the right
                if position != omni.ui.DockPosition.SAME:
                    await omni.kit.app.get_app().next_update_async()
                    await omni.kit.app.get_app().next_update_async()
                    window.dock_in(dockspace, position, ratio)

                # This genrally works in a variety of cases from load, re-load, and save extesnion .py file reload
                # But it depends on a call order in omni.ui and registered selected_in_dock and dock_changed callbacks.
                await omni.kit.app.get_app().next_update_async()
                updates_enabled = window.selected_in_dock if window.docked else window.visible
                window.viewport_api.updates_enabled = updates_enabled
                # Set the dock-tab visible to the state of the default Viewport
                window.dock_tab_bar_visible = dock_tab_bar_visible
                dockspace.dock_tab_bar_visible = dock_tab_bar_visible

        import asyncio
        if self.__window_name == 'Viewport':
            asyncio.ensure_future(wait_for_window(self.__window_name, 'Viewport', omni.ui.DockPosition.SAME))
        else:
            asyncio.ensure_future(wait_for_window(self.__window_name, 'Viewport', omni.ui.DockPosition.RIGHT, 0.5))

    def __show_window(self, menu, visible):
        if visible:
            if not self.__window:
                def visiblity_changed(visible):
                    if not visible:
                        self.__show_window(None, False)
                self.__window = ViewportWindow(self.__window_name)
                self.__window.set_visibility_changed_fn(visiblity_changed)
                self.__dock_window()
            self.__window.visible = True
        elif self.__window:
            # For now just hide the Window. May want to add a setting for a full destroy-rebuild cycle, but now
            # that won't work well with the get_frame API.
            if True:
                self.__window.visible = False
            else:
                self.__window.set_visibility_changed_fn(None)
                self.__window.destroy()
                self.__window = None

        omni.kit.menu.utils.refresh_menu_items("Window")

    def __is_window_visible(self) -> bool:
        return False if self.__window is None else self.__window.visible


class ViewportWindowExtension(omni.ext.IExt):
    """The Entry Point for Viewport Window
    """

    def __init__(self):
        self.__registered = None
        self.__vp_items = None
        super().__init__()
        self.__extension_enabled_hooks = []
        self.__extension_dependent_scenes = {}

    def on_startup(self, extension_id: str):
        settings = carb.settings.get_settings()
        default_name = settings.get(DEFAULT_VIEWPORT_NAME) or 'Viewport'
        self.__vp_items = (
            _MenuEntry(default_name, not bool(settings.get(DEFAULT_VIEWPORT_NO_OPEN)), 0),
        )
        if True:
            self.__vp_items = (
                self.__vp_items[0],
                _MenuEntry(f'{default_name}', False, 1),
            )

        self.__registered = self.__register_scenes()
        self.__default_drag_handlers = (
            UsdFileDropDelegate('/persistent/app/viewport/previewOnPeek'),
            UsdShadeDropDelegate(),
            MaterialFileDropDelegate(),
            AudioFileDropDelegate()
        )

    def on_shutdown(self):
        self.__extension_enabled_hooks = []

        for vp_entry in (self.__vp_items or ()):
            vp_entry.destroy()
        self.__vp_items = None

        self.__default_drag_handlers = None

        if self.__extension_dependent_scenes:
            self.__unregister_scenes(self.__extension_dependent_scenes.values())
            self.__extension_dependent_scenes = None

        if self.__registered:
            self.__unregister_scenes(self.__registered)
            self.__registered = None

        from .events import set_ui_delegate
        set_ui_delegate(None)

    def __register_ext_dependent_scene(self, ext_name: str, loaded_exts: List[Dict], ext_manager: omni.ext.ExtensionManager,
                                       factory_dict: Dict, other_factories: Optional[Dict] = None):
        def is_extension_loaded(ext_name: str) -> bool:
            for ext in loaded_exts:
                if ext_name == omni.ext.get_extension_name(ext['id']):
                    return ext['enabled']

        def extension_loaded(ext_id: str, log_msg_str: str = "Removing"):
            nonlocal ext_name
            # UnRegister existing first
            if other_factories:
                carb.log_info(f'{log_msg_str} embedded ui.scenes fallbacks for "{ext_name}"')
                for factory_id, factory in other_factories.items():
                    self.__extension_dependent_scenes[factory_id] = None

            # Register new scene types second
            carb.log_info(f'Adding ui.scenes dependent on "{ext_name}"')
            from omni.kit.viewport.registry import RegisterScene
            self.__extension_dependent_scenes[ext_name] = [
                RegisterScene(factory, factory_id) for factory_id, factory in factory_dict.items()
            ]

        def extension_unloaded(ext_id: str):
            nonlocal ext_name
            # UnRegister scenes dependeing on extension first
            if self.__extension_dependent_scenes.get(ext_name):
                carb.log_info(f'Removing ui.scenes dependent on "{ext_name}"')
                self.__extension_dependent_scenes[ext_name] = None

            # Register replacement scene types second
            if other_factories:
                carb.log_info(f'Adding embedded ui.scenes fallbacks for "{ext_name}"')
                from omni.kit.viewport.registry import RegisterScene
                for factory_id, factory in other_factories.items():
                    self.__extension_dependent_scenes[factory_id] = RegisterScene(factory, factory_id)

        if not is_extension_loaded(ext_name):
            carb.log_info(f'{ext_name} is not loaded, adding extension enabled hooks')
            self.__extension_enabled_hooks += [
                ext_manager.subscribe_to_extension_enable(
                    extension_loaded,
                    extension_unloaded,
                    ext_name=ext_name,
                    hook_name="omni.kit.viewport.window-ext_name"
                )
            ]
            if other_factories:
                extension_unloaded(ext_name)
        else:
            extension_loaded(ext_name, "Ignoring")

    def __register_scenes(self):
        # Register all of the items that use omni.ui.scene to add functionality
        from omni.kit.viewport.registry import RegisterScene, RegisterViewportLayer

        app = omni.kit.app.get_app_interface()
        ext_mgr = app.get_extension_manager()
        exts_loaded = ext_mgr.get_extensions()
        registered = []

        # Register the Camera manipulator (if available)
        def delay_load_cam_manip(*args, **kwargs):
            from .manipulator.camera import ViewportCameraManiulatorFactory
            return ViewportCameraManiulatorFactory(*args, **kwargs)

        self.__register_ext_dependent_scene('omni.kit.manipulator.camera', exts_loaded, ext_mgr, {
            'omni.kit.viewport.window.manipulator.Camera': delay_load_cam_manip,
        })

        # Register the Selection manipulator (if available)
        def delay_load_selection_manip(*args, **kwargs):
            from .manipulator.selection import SelectionManipulatorItem
            return SelectionManipulatorItem(*args, **kwargs)

        self.__register_ext_dependent_scene('omni.kit.manipulator.selection', exts_loaded, ext_mgr, {
            'omni.kit.viewport.window.manipulator.Selection': delay_load_selection_manip,
        })

        # Register the legacy Gizmo drawing (if available) or the omni.ui.scene version when not
        from .scene.legacy import LegacyGridScene, LegacyLightScene, LegacyAudioScene
        from .scene.scenes import SimpleGrid
        self.__register_ext_dependent_scene("omni.kit.viewport.legacy_gizmos", exts_loaded, ext_mgr, {
            'omni.kit.viewport.window.scene.LegacyGrid': LegacyGridScene,
            'omni.kit.viewport.window.scene.LegacyLight': LegacyLightScene,
            'omni.kit.viewport.window.scene.LegacyAudio': LegacyAudioScene,
        }, {
            'omni.kit.viewport.window.scene.SimpleGrid': SimpleGrid
        })

        # Register other omni.scene.ui elements
        from .scene.scenes import SimpleGrid, SimpleOrigin, CameraAxisLayer
        registered = [
            RegisterScene(SimpleOrigin, 'omni.kit.viewport.window.scene.SimpleOrigin'),
            RegisterViewportLayer(CameraAxisLayer, 'omni.kit.viewport.window.CameraAxisLayer')
        ]

        # Register the context click menu
        from .manipulator.context_menu import ViewportClickFactory
        registered += [
            RegisterScene(ViewportClickFactory, 'omni.kit.viewport.window.manipulator.ContextMenu')
        ]

        # Register the Object click manipulator
        from .manipulator.object_click import ObjectClickFactory
        registered += [
            RegisterScene(ObjectClickFactory, 'omni.kit.viewport.window.manipulator.ObjectClick')
        ]

        # Register the HUD stats
        from .stats import ViewportStatsLayer
        registered.append(RegisterViewportLayer(ViewportStatsLayer, 'omni.kit.viewport.window.ViewportStats'))

        # Finally register the ViewportSceneLayer
        from .scene.layer import ViewportSceneLayer
        registered.append(RegisterViewportLayer(ViewportSceneLayer, 'omni.kit.viewport.window.SceneLayer'))
        return registered

    def __unregister_scenes(self, registered):
        for item in registered:
            try:
                item.destroy()
            except Exception:
                pass
