# Copyright (c) 2018-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.
#
from typing import Callable
import carb
import os
import asyncio
import omni.ext
import omni.usd
import omni.kit.ui
import omni.kit.app
import omni.ui as ui
import omni.kit.notification_manager as nm
import omni.kit.usd.layers as layers

from functools import partial
from omni.kit.usd.layers import LayerUtils
from typing import List
from pxr import Sdf

from .window import LayerWindow
from .layer_item import LayerItem
from .layer_model import LayerModel
from .layer_link_window import LayerLinkWindow
from .layer_model_utils import LayerModelUtils
from .actions import ActionManager


_extension_instance = None


class LayerExtension(omni.ext.IExt):
    """The entry point for Layer 2"""

    CONTEXT_MENU_ITEM_INSERT_SUBLAYER = "Insert As Sublayer"
    WINDOW_NAME = "Layer"
    MENU_PATH = f"Window/{WINDOW_NAME}"

    def on_startup(self, ext_id):
        global _extension_instance
        _extension_instance = self

        # Register command related actions
        self._ext_name = omni.ext.get_extension_name(ext_id)
        self._action_manager = ActionManager()
        self._action_manager.on_startup(self)

        ui.Workspace.set_show_window_fn(LayerExtension.WINDOW_NAME, partial(self.show_window, None))

        self._usd_context = omni.usd.get_context()
        self._layers = layers.get_layers(self._usd_context)
        self._selection_listeners = set([])
        self._window = None
        try:
            self._menu = omni.kit.ui.get_editor_menu().add_item(
                LayerExtension.MENU_PATH, self.show_window, toggle=True, value=True
            )
        except Exception:
            self._menu = None
        ui.Workspace.show_window(LayerExtension.WINDOW_NAME)

        self._link_window = None

        app = omni.kit.app.get_app()
        manager = app.get_extension_manager()
        self._context_icon_menu_items = []
        self._extensions_subscription = manager.get_change_event_stream().create_subscription_to_pop(
            self._on_event, name="omni.kit.widget.layers"
        )
        self._menu_registered = False

        self._on_layer_edit_mode_update(self._layers.get_edit_mode())
        self._layer_edit_mode_subscription = self._layers.get_event_stream().create_subscription_to_pop_by_type(
            int(layers.LayerEventType.EDIT_MODE_CHANGED),
            self._on_layer_events, name="Layers Extension"
        )

    def _on_layer_events(self, event: carb.events.IEvent):
        payload = layers.get_layer_event_payload(event)
        if payload and payload.event_type == layers.LayerEventType.EDIT_MODE_CHANGED:
            edit_mode = self._layers.get_edit_mode()
            self._on_layer_edit_mode_update(edit_mode)

    def _on_layer_edit_mode_update(self, edit_mode):
        if edit_mode == layers.LayerEditMode.SPECS_LINKING:
            if not self._link_window:
                self._link_window = LayerLinkWindow()
        elif self._link_window:
            self._link_window.destroy()
            self._link_window = None

    def _on_event(self, event):
        if not self._menu_registered:
            self._register_menus()
            self._menu_registered = True

    def on_shutdown(self):
        self._action_manager.on_shutdown()

        # Clean up global resources
        LayerModelUtils.on_shutdown()

        global _extension_instance
        _extension_instance = None
        if self._link_window:
            self._link_window.destroy()
            self._link_window = None

        # remove menu
        self._menu = None
        if self._window:
            self._window.destroy()
            self._window = None

        for fn in self._selection_listeners:
            fn(None)
        self._selection_listeners.clear()

        self._menu_registered = False
        self._extensions_subscription = None
        self._unregister_menus()
        self._layers = None

        ui.Workspace.set_show_window_fn(LayerExtension.WINDOW_NAME, None)

    def _get_content_window(self):
        try:
            import omni.kit.window.content_browser as content

            return content.get_content_window()
        except ImportError:
            pass

        return None

    def _unregister_menus(self):
        self._context_icon_menu_items.clear()
        extension_manager = omni.kit.app.get_app_interface().get_extension_manager()

        content_window = self._get_content_window()
        if content_window and extension_manager.is_extension_enabled("omni.kit.window.content_browser"):
            content_window.delete_context_menu(self.CONTEXT_MENU_ITEM_INSERT_SUBLAYER)

    def _register_menus(self):
        content_window = self._get_content_window()

        if content_window:
            content_window.add_context_menu(
                self.CONTEXT_MENU_ITEM_INSERT_SUBLAYER,
                "external_link.svg",
                lambda b, c: self._on_icon_menu_click(b, c),
                LayerExtension._is_show_insert_visible,
                separator_name=None,
                index=3  # Index
            )

    @staticmethod
    def _is_show_insert_visible(content_url):
        return omni.usd.is_usd_writable_filetype(content_url)

    def _on_icon_menu_click(self, menu, value):
        file_path = value
        stage = self._usd_context.get_stage()
        if not stage:
            return

        root_layer = stage.GetRootLayer()
        if root_layer.identifier != file_path:
            found = False
            for sublayer_path in root_layer.subLayerPaths:
                absolute_path = root_layer.ComputeAbsolutePath(sublayer_path)
                if os.path.normpath(value) == os.path.normpath(absolute_path):
                    found = True

            if found:
                nm.post_notification(
                    f"Duplicate sublayer found in the Root Layer.",
                    status=nm.NotificationStatus.WARNING,
                    duration=4
                )
            else:
                LayerUtils.insert_sublayer(root_layer, 0, file_path)
        else:
            nm.post_notification(
                f"Cannot insert Root Layer as sublayer.",
                status=nm.NotificationStatus.WARNING,
                duration=4
            )

    async def _destroy_window_async(self):
        # wait one frame, this is due to the one frame defer
        # in Window::_moveToMainOSWindow()
        await omni.kit.app.get_app().next_update_async()
        if self._window:
            self._window.destroy()
            self._window = None

    def _visiblity_changed_fn(self, visible):
        if self._menu:
            omni.kit.ui.get_editor_menu().set_value(LayerExtension.MENU_PATH, visible)
        if not visible:
            # Destroy the window, since we are creating new window
            # in show_window
            asyncio.ensure_future(self._destroy_window_async())

    def show_window(self, menu, value):
        if value and not self._window:
            self._window = LayerWindow(LayerExtension.WINDOW_NAME, self._usd_context)
            self._window.set_visibility_changed_listener(self._visiblity_changed_fn)
            # Register selection listeners again
            for fn in self._selection_listeners:
                self._window.add_layer_selection_changed_fn(fn)
        elif self._window:
            self._window.set_visible(value)
            for fn in self._selection_listeners:
                fn(None)

    def get_current_focused_layer_item(self) -> LayerItem:
        """Gets the current focused layer item in Layer Window."""
        if not self._window:
            return None

        return self._window.get_current_focused_layer_item()

    def set_current_focused_layer_item(self, layer_identifier: str):
        """Sets the focused layer item in Layer Window."""
        if not self._window:
            return

        return self._window.set_current_focused_layer_item(layer_identifier)

    def remove_layer_selection_changed_fn(self, fn: Callable[[LayerItem], None]):
        """Removes selection listener."""
        if not fn:
            return

        self._selection_listeners.discard(fn)

        if not self._window:
            return

        return self._window.remove_layer_selection_changed_fn(fn)

    def add_layer_selection_changed_fn(self, fn: Callable[[LayerItem], None]):
        if not fn:
            return

        self._selection_listeners.add(fn)

        if not self._window:
            return

        return self._window.add_layer_selection_changed_fn(fn)

    def get_layer_model(self) -> LayerModel:
        if not self._window:
            return None

        return self._window.get_layer_model()

    def get_selected_items(self) -> List[ui.AbstractItem]:
        if not self._window:
            return []

        return list(self._window.layer_view.selection)

    @staticmethod
    def get_instance():
        return _extension_instance
