## Copyright (c) 2021, 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 urllib.parse
import carb.input
import omni.client
import omni.ext
import omni.kit.menu.utils
import omni.kit.ui
import omni.usd

from omni.kit.helper.file_utils import get_latest_urls_from_event_queue, asset_types, FILE_EVENT_QUEUE_UPDATED
from omni.kit.menu.utils import MenuItemDescription
from .file_actions import register_actions, deregister_actions
from omni.ui import color as cl
from pathlib import Path

_extension_instance = None
_extension_path = None
INTERACTIVE_TEXT = cl.shade(cl("#1A91C5"))
SHARE_ICON_PATH = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)).joinpath("data/icons/share.svg")

class FileMenuExtension(omni.ext.IExt):
    def __init__(self):
        super().__init__()
        omni.kit.menu.utils.set_default_menu_proirity("File", -10)

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

        global _extension_path
        _extension_path = omni.kit.app.get_app_interface().get_extension_manager().get_extension_path(ext_id)

        self._ext_name = omni.ext.get_extension_name(ext_id)
        register_actions(self._ext_name)

        settings = carb.settings.get_settings()
        self._max_recent_files = settings.get("exts/omni.kit.menu.file/maxRecentFiles") or 10

        self._file_menu_list = None
        self._recent_menu_list = None
        self._build_file_menu()

        event_stream = omni.kit.app.get_app().get_message_bus_event_stream()
        self._event_sub = event_stream.create_subscription_to_pop_by_type(FILE_EVENT_QUEUE_UPDATED, lambda _: self._build_recent_menu())

        events = omni.usd.get_context().get_stage_event_stream()
        self._stage_event_subscription = events.create_subscription_to_pop(self._on_stage_event, name="omni.kit.menu.file stage watcher")

    def on_shutdown(self):
        global _extension_instance

        deregister_actions(self._ext_name)

        _extension_instance = None
        self._stage_sub = None
        omni.kit.menu.utils.remove_menu_items(self._recent_menu_list, "File")
        omni.kit.menu.utils.remove_menu_items(self._file_menu_list, "File")
        self._recent_menu_list = None
        self._file_menu_list = None
        self._event_sub = None

    def _on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.CLOSED):
            self._build_file_menu()

    def _build_file_menu(self):
        # setup menu
        self._file_menu_list = [
            MenuItemDescription(
                name="New",
                glyph="file.svg",
                onclick_action=("omni.kit.window.file", "new"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.N),
            ),
            MenuItemDescription(
                name="Open",
                glyph="folder_open.svg",
                onclick_action=("omni.kit.window.file", "open"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.O),
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Re-open with New Edit Layer",
                glyph="external_link.svg",
                # enable_fn=lambda: not FileMenuExtension.is_new_stage(),
                onclick_action=("omni.kit.menu.file", "open_stage_with_new_edit_layer")
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Share",
                glyph=str(SHARE_ICON_PATH),
                enable_fn=FileMenuExtension.can_share,
                onclick_action=("omni.kit.window.file", "share")
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Save",
                glyph="save.svg",
                # enable_fn=FileMenuExtension.can_close,
                onclick_action=("omni.kit.window.file", "save"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.S),
            ),
            MenuItemDescription(
                name="Save With Options",
                glyph="save.svg",
                # enable_fn=FileMenuExtension.can_close,
                onclick_action=("omni.kit.window.file", "save_with_options"),
                hotkey=(
                    carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL | carb.input.KEYBOARD_MODIFIER_FLAG_ALT,
                    carb.input.KeyboardInput.S,
                ),
            ),
            MenuItemDescription(
                name="Save As...",
                glyph="none.svg",
                # enable_fn=FileMenuExtension.can_close,
                onclick_action=("omni.kit.window.file", "save_as"),
                hotkey=(
                    carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL | carb.input.KEYBOARD_MODIFIER_FLAG_SHIFT,
                    carb.input.KeyboardInput.S,
                ),
            ),
            MenuItemDescription(
                name="Save Flattened As...",
                glyph="none.svg",
                # enable_fn=FileMenuExtension.can_close,
                onclick_action=("omni.kit.window.file", "save_as_flattened"),
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Add Reference",
                glyph="none.svg",
                onclick_action=("omni.kit.window.file", "add_reference"),
            ),
            MenuItemDescription(
                name="Add Payload",
                glyph="none.svg",
                onclick_action=("omni.kit.window.file", "add_payload")
            ),
            MenuItemDescription(),
            MenuItemDescription(name="Exit", glyph="none.svg", onclick_action=("omni.kit.menu.file", "quit")),
        ]
        if not carb.settings.get_settings().get("/app/fastShutdown"):
            self._file_menu_list.append(MenuItemDescription(name="Exit Fast (Experimental)", glyph="none.svg", onclick_action=("omni.kit.menu.file", "quit_fast")))

        self._build_recent_menu()
        omni.kit.menu.utils.add_menu_items(self._file_menu_list, "File", -10)

    def _build_recent_menu(self):
        if self._recent_menu_list:
            omni.kit.menu.utils.remove_menu_items(self._recent_menu_list, "File", rebuild_menus=False)

        recent_files = get_latest_urls_from_event_queue(self._max_recent_files, asset_type=asset_types.ASSET_TYPE_USD)
        sub_menu = []
        # add reopen
        stage = omni.usd.get_context().get_stage()
        is_anonymous = stage.GetRootLayer().anonymous if stage else False
        filename_url = ""
        if not is_anonymous:
            filename_url = stage.GetRootLayer().identifier if stage else ""

        if filename_url:
            sub_menu.append(
                MenuItemDescription(
                    name=f"Current stage: {filename_url}",
                    onclick_action=("omni.kit.window.file", "reopen"),
                    user={"user_style": {"color": INTERACTIVE_TEXT}}
                    )
                )
        elif not recent_files:
            sub_menu.append(MenuItemDescription(name="None", enabled=False))

        if recent_files:
            for recent_file in recent_files:
                # NOTE: not compatible with hotkeys as passing URL to open_stage
                sub_menu.append(
                    MenuItemDescription(
                        name=urllib.parse.unquote(recent_file),
                        onclick_action=("omni.kit.window.file", "open_stage", recent_file),
                    )
                )

        self._recent_menu_list = [
            MenuItemDescription(name="Open Recent", glyph="none.svg", appear_after="Open", sub_menu=sub_menu)
        ]
        omni.kit.menu.utils.add_menu_items(self._recent_menu_list, "File")

    def is_new_stage():
        return omni.usd.get_context().is_new_stage()

    def can_open():
        stage_state = omni.usd.get_context().get_stage_state()
        return stage_state == omni.usd.StageState.OPENED or stage_state == omni.usd.StageState.CLOSED

    def can_save():
        return (
            omni.usd.get_context().get_stage_state() == omni.usd.StageState.OPENED
            and not FileMenuExtension.is_new_stage()
            and omni.usd.get_context().is_writable()
        )

    def can_close():
        return omni.usd.get_context().get_stage_state() == omni.usd.StageState.OPENED

    def can_close_and_not_is_new_stage():
        return FileMenuExtension.can_close() and not FileMenuExtension.is_new_stage()

    def can_share():
        return (
            omni.usd.get_context() is not None
            and omni.usd.get_context().get_stage_state() == omni.usd.StageState.OPENED
            and omni.usd.get_context().get_stage_url().startswith("omniverse://")
        )

def get_extension_path(sub_directory):
    global _extension_path
    path = _extension_path
    if sub_directory:
        path = os.path.normpath(os.path.join(path, sub_directory))
    return path


def get_instance():
    global _extension_instance
    return _extension_instance
