# 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.
#
import asyncio
import os
import omni.ui as ui

from functools import partial
from typing import Callable
from omni.kit.widget.filebrowser import FileBrowserItem
from omni.kit.widget.versioning import CheckpointWidget
from .view import FilePickerView
from .file_ops import *
from .style import get_style


class BaseContextMenu:
    """
    Base class popup menu for the hovered FileBrowserItem.  Provides API for users to add menu items.

    """
    def __init__(self, title: str = None, **kwargs):
        self._title: str = title
        self._view: FilePickerView = kwargs.get("view", None)
        self._checkpoint: CheckpointWidget = kwargs.get("checkpoint", None)
        self._context_menu: ui.Menu = None
        self._menu_dict: list = []
        self._menu_items: list = []
        self._context: dict = None

    @property
    def menu(self) -> ui.Menu:
        """:obj:`omni.ui.Menu` The menu widget"""

        return self._context_menu

    @property
    def context(self) -> dict:
        """dict: Provides data to the callback.  Available keys are {'item', 'selected'}"""
        return self._context

    def show(
        self,
        item: FileBrowserItem,
        selected: List[FileBrowserItem] = [],
    ):
        """
        Creates the popup menu from definition for immediate display. Receives as input, information about the
        item.  These values are made available to the callback via the 'context' dictionary.

        Args:
            item (FileBrowseritem): Item for which to create menu.,
            selected ([FileBrowserItem]): List of currently selected items. Default [].

        """
        self._context = {}
        self._context["item"] = item
        self._context["selected"] = selected

        self._context_menu = ui.Menu(self._title, style=get_style(), style_type_name_override="Menu")
        menu_entry_built = False
        with self._context_menu:
            prev_entry = None
            for i, menu_entry in enumerate(self._menu_dict):
                if i == len(self._menu_dict) - 1 and not menu_entry.get("name"):
                    # Don't draw separator as last item
                    break
                if menu_entry.get("name") or (prev_entry and prev_entry.get("name")):
                    # Don't draw 2 separators in a row
                    entry_built = self._build_menu(menu_entry, self._context)
                    if entry_built:
                        # only proceed prev_entry if the current entry is built, so that if we have menu items that
                        # are not built in between 2 separators, the 2nd separator will not build
                        prev_entry = menu_entry
                    menu_entry_built |= entry_built

        # Show it (if there's at least one menu entry built)
        if len(self._menu_dict) > 0 and menu_entry_built:
            self._context_menu.show()

    def hide(self):
        self._context_menu.hide()

    def add_menu_item(self, name: str, glyph: str, onclick_fn: Callable, show_fn: Callable, index: int = -1, separator_name: Optional[str] = None) -> str:
        """
        Adds menu item, with corresponding callbacks, to this context menu.

        Args:
            name (str): Name of the menu item (e.g. 'Open'), this name must be unique across the context menu.
            glyph (str): Associated glyph to display for this menu item.
            onclick_fn (Callable): This callback function is executed when the menu item is clicked. Function signature:
                void fn(context: Dict)
            show_fn (Callable): Returns True to display this menu item. Function signature: bool fn(context: Dict).
                For example, test filename extension to decide whether to display a 'Play Sound' action.
            index (int): The position that this menu item will be inserted to.
            separator_name (str): The separator name of the separator menu item. Default to '_placeholder_'. When the
                index is not explicitly set, or if the index is out of range, this will be used to locate where to add
                the menu item; if specified, the index passed in will be counted from the saparator with the provided
                name. This is for OM-86768 as part of the effort to match Navigator and Kit UX for Filepicker/Content Browser for context menus.

        Returns:
            str: Name of menu item if successful, None otherwise.

        """
        if not name:
            return None
        elif name in [item.get("name", None) for item in self._menu_dict]:
            # Reject duplicates
            return None

        menu_item = {"name": name, "glyph": glyph or ""}
        if onclick_fn:
            menu_item["onclick_fn"] = lambda context, name=name: onclick_fn(name, context["item"].path)
        if show_fn:
            menu_item["show_fn"] = lambda context: show_fn(context["item"].path)

        if index >= 0 and index <= len(self._menu_dict) and separator_name is None:
            pass
        else:
            if separator_name is None:
                separator_name = "_placeholder_"
            placeholder_index = (i for i, item in enumerate(self._menu_dict) if separator_name in item)
            matched = next(placeholder_index, len(self._menu_dict))
            index = max(min(matched + index, len(self._menu_dict)), 0)

        self._menu_dict.insert(index, menu_item)

        return name

    def delete_menu_item(self, name: str):
        """
        Deletes the menu item, with the given name, from this context menu.

        Args:
            name (str): Name of the menu item (e.g. 'Open').

        """
        if not name:
            return
        found = (i for i, item in enumerate(self._menu_dict) if name == item.get("name", None))
        for j in sorted([i for i in found], reverse=True):
            del self._menu_dict[j]

    def _build_menu(self, menu_entry, context):
        menu_entry_built = False
        from omni.kit.ui import get_custom_glyph_code

        if "name" in menu_entry and isinstance(menu_entry["name"], dict):
            sub_menu = copy.copy(menu_entry["name"])
            icon = "plus.svg"
            if "glyph" in sub_menu:
                icon = sub_menu["glyph"]
                del sub_menu["glyph"]

            gylph_char = get_custom_glyph_code("${glyphs}/" + icon)
            for item in sub_menu:
                menu_item = ui.Menu(f" {gylph_char}   {item}")
                with menu_item:
                    for menu_entry in sub_menu[item]:
                        self._build_menu(menu_entry, context)
                    menu_entry_built = True
                self._menu_items.append(menu_item)
            return menu_entry_built

        if not self._get_fn_result(menu_entry, "show_fn", context):
            return False

        if "populate_fn" in menu_entry:
            menu_item = menu_entry["populate_fn"](context)
            self._menu_items.append(menu_item)
            return True

        if menu_entry["name"] == "":
            menu_item = ui.Separator(style_type_name_override="Menu.Separator")
        else:
            menu_name = ""
            if "glyph" in menu_entry:
                menu_name = " " + get_custom_glyph_code("${glyphs}/" + menu_entry["glyph"]) + "   "

            if "onclick_fn" in menu_entry and menu_entry["onclick_fn"] is not None:
                enabled = self._get_fn_result(menu_entry, "enabled_fn", context)
                menu_item = ui.MenuItem(
                    menu_name + menu_entry["name"],
                    triggered_fn=partial(menu_entry["onclick_fn"], context),
                    enabled=enabled,
                    style_type_name_override="Menu.Item",
                )
            else:
                menu_item = ui.MenuItem(menu_name + menu_entry["name"], enabled=False, style_type_name_override="Menu.Item")

            if "show_fn_async" in menu_entry:
                menu_item.visible = False

                async def _delay_update_menu_item():
                    await menu_entry["show_fn_async"](context, menu_item)
                    # update menu item visiblity after the show fn result is returned
                    self._update_menu_item_visibility()

                asyncio.ensure_future(_delay_update_menu_item())
        # FIXME: in the case of when there's only one menu entry and it's show_fn_async returned false, this result will
        # be wrong
        self._menu_items.append(menu_item)
        return True

    def _get_fn_result(self, menu_entry, name, context):
        if not name in menu_entry:
            return True
        if menu_entry[name] is None:
            return True
        if isinstance(menu_entry[name], list):
            for show_fn in menu_entry[name]:
                if not show_fn(context):
                    return False
        else:
            if not menu_entry[name](context):
                return False
        return True

    def _update_menu_item_visibility(self):
        """
        Utility to update menu item visiblity.
        Mainly to eliminate the case where 2 separators are shown next to each other because of the menu items in
        between them are async and made the menu items invisible.
        """
        if not self._context_menu:
            return

        last_visible_item = None
        for item in self._menu_items:
            if not item.visible:
                continue
            # if the current item is a separator, check that the previous visible menu item is not separator to avoid
            # having 2 continuous separators
            if isinstance(item, ui.Separator) and last_visible_item and isinstance(last_visible_item, ui.Separator) and\
                last_visible_item.visible:
                    item.visible = False
            last_visible_item = item
        self._context_menu.show()

    def destroy(self):
        self._context_menu = None
        self._menu_dict = None
        self._context = None
        self._menu_items.clear()


class ContextMenu(BaseContextMenu):
    """
    Creates popup menu for the hovered FileBrowserItem.  In addition to the set of default actions below,
    users can add more via the add_menu_item API.

    """
    def __init__(self, **kwargs):
        super().__init__(title="Context menu", **kwargs)

        self._menu_dict = [
            {
                "name": "Open in File Browser",
                "glyph": "folder_open.svg",
                "onclick_fn": lambda context: open_in_file_browser(context["item"]),
                "show_fn": [lambda context: os.path.exists(context["item"].path)],
            },
            {
                "name": "New USD File",
                "glyph": "file.svg",
                "onclick_fn": lambda context: create_usd_file(context["item"]),
                "show_fn": [
                    lambda context: context["item"].is_folder,
                    lambda _: is_usd_supported(),
                    # OM-72882: should not allow creating folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: len(context["selected"]) <= 1,
                    lambda context: not context["item"] in self._view.all_collection_items(collection="omniverse"),
                ]
            },
            {
                "name": "Restore",
                "glyph": "restore.svg",
                "onclick_fn": lambda context: restore_items(context["selected"], self._view),
                "show_fn": [
                    # OM-72882: should not allow deleting folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: context["item"].is_deleted,
                ],
            },
            {
                "name": "Refresh",
                "glyph": "menu_refresh.svg",
                "onclick_fn": lambda context: refresh_item(context["item"], self._view),
                "show_fn": [lambda context: len(context["selected"]) <= 1,]
            },
            {
                "name": "",
                "_add_on_begin_separator_": "",
            },
            {
                "name": "",
                "_add_on_end_separator_": "",
            },
            {
                "name": "New Folder",
                "glyph": "menu_plus.svg",
                "onclick_fn": lambda context: create_folder(context["item"]),
                "show_fn": [
                    lambda context: context["item"].is_folder,
                    # OM-72882: should not allow creating folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: len(context["selected"]) <= 1,
                    # TODO: Can we create any folder in the omniverer server's root folder?
                    #lambda context: not(context["item"].parent and \
                    #         context["item"].parent.path == "omniverse://"),
                ],
            },
            {
                "name": "",
            },
            {
                "name": "Create Checkpoint",
                "glyph": "database.svg",
                "onclick_fn": lambda context: checkpoint_items(context["selected"], self._checkpoint),
                "show_fn": [
                    lambda context: not context["item"].is_folder,
                    lambda context: len(context["selected"]) == 1,
                ],
                "show_fn_async": is_item_checkpointable,
            },
            {
                "name": "",
            },
            {
                "name": "Rename",
                "glyph": "pencil.svg",
                "onclick_fn": lambda context: rename_item(context["item"], self._view),
                "show_fn": [
                    lambda context: len(context["selected"]) == 1,
                    lambda context: context["item"].writeable,
                ]
            },
            {
                "name": "Delete",
                "glyph": "menu_delete.svg",
                "onclick_fn": lambda context: delete_items(context["selected"]),
                "show_fn": [
                    # OM-72882: should not allow deleting folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: not context["item"].is_deleted,
                    # OM-94626: we couldn't delete when select nothing
                    lambda context: len(context["selected"]) >= 1,
                ],
            },
            {
                "name": "",
                "_placeholder_": ""
            },
            {
                "name": "Add Bookmark",
                "glyph": "bookmark.svg",
                "onclick_fn": lambda context: add_bookmark(context["item"], self._view),
                "show_fn": [lambda context: len(context["selected"]) <= 1,]
            },
            {
                "name": "Copy URL Link",
                "glyph": "share.svg",
                "onclick_fn": lambda context: copy_to_clipboard(context["item"]),
                "show_fn": [lambda context: len(context["selected"]) <= 1,]
            },
            {
                "name": "Obliterate",
                "glyph": "menu_delete.svg",
                "onclick_fn": lambda context: obliterate_items(context["selected"], self._view),
                "show_fn": [
                    # OM-72882: should not allow deleting folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: context["item"].is_deleted,
                ],
            },
        ]


class UdimContextMenu(BaseContextMenu):
    """Creates popup menu for the hovered FileBrowserItem that are Udim nodes."""
    def __init__(self, **kwargs):
        super().__init__(title="Udim menu", **kwargs)

        self._menu_dict = [
            {
                "name": "Add Bookmark",
                "glyph": "bookmark.svg",
                "onclick_fn": lambda context: add_bookmark(context["item"], self._view),
            },
        ]


class CollectionContextMenu(BaseContextMenu):
    """Creates popup menu for the hovered FileBrowserItem that are collection nodes."""
    def __init__(self, **kwargs):
        super().__init__(title="Collection menu", **kwargs)

        # OM-75883: override menu dict for "Omniverse" collection node, there should be only
        # one menu entry, matching Navigator UX
        self._menu_dict = [
            {
                "name": "Add Server",
                "glyph": "menu_plus.svg",
                "onclick_fn": lambda context: add_connection(self._view),
                "show_fn": [
                    lambda context: context['item'].name == "Omniverse",
                    lambda context: carb.settings.get_settings().get_as_bool("/exts/omni.kit.window.filepicker/show_add_new_connection"),
                ],
            },
        ]


class ConnectionContextMenu(BaseContextMenu):
    """Creates popup menu for the server connection FileBrowserItem grouped under Omniverse collection node."""
    def __init__(self, **kwargs):
        super().__init__(title="Connection menu", **kwargs)

        # OM-86768: matching context menu with Navigator UX for server connection
        # TODO: Add API Tokens and Clear Cached Folders once API is exposed in omni.client
        self._menu_dict = [
            {
                "name": "New Folder",
                "glyph": "menu_plus.svg",
                "onclick_fn": lambda context: create_folder(context["item"]),
                "show_fn": [
                    # OM-72882: should not allow creating folder/file in read-only directory
                    lambda context: context["item"].writeable,
                    lambda context: len(context["selected"]) == 1,
                ],
            },
            {
                "name": "Reconnect Server",
                "glyph": "menu_refresh.svg",
                "onclick_fn": lambda context: refresh_connection(context["item"], self._view),
            },
            {
                "name": "",
            },
            {
                "name": "Rename",
                "glyph": "pencil.svg",
                "onclick_fn": lambda context: rename_item(context["item"], self._view),
                "show_fn": [
                    lambda context: len(context["selected"]) == 1,
                ]
            },
            {
                "name": "",
            },
            {
                "name": "Log In",
                "glyph": "user.svg",
                "onclick_fn": lambda context: refresh_connection(context["item"], self._view),
                "show_fn": [
                    lambda context: not FilePickerView.is_connected(context["item"].path),
                ],
            },
            {
                "name": "Log Out",
                "glyph": "user.svg",
                "onclick_fn": lambda context: log_out_from_connection(context["item"], self._view),
                "show_fn": [
                    lambda context: FilePickerView.is_connected(context["item"].path),
                ],
            },
            {
                "name": "",
            },
            {
                "name": "About",
                "glyph": "question.svg",
                "onclick_fn": lambda context: about_connection(context["item"]),
            },
            {
                "name": "Remove Server",
                "glyph": "menu_delete.svg",
                "onclick_fn": lambda context: remove_connection(context["item"], self._view),
            },
        ]


class BookmarkContextMenu(BaseContextMenu):
    """Creates popup menu for BookmarkItems."""
    def __init__(self, **kwargs):
        super().__init__(title="Bookmark menu", **kwargs)

        # OM-66726: Edit bookmark similar to Navigator
        self._menu_dict = [
            {
                "name": "Edit",
                "glyph": "pencil.svg",
                "onclick_fn": lambda context: edit_bookmark(context['item'], self._view),
            },
            {
                "name": "Delete",
                "glyph": "menu_delete.svg",
                "onclick_fn": lambda context: delete_bookmark(context['item'], self._view),
            },
        ]
