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

import os.path
import typing
from functools import partial

import omni.ext
import omni.kit.app
import omni.kit.ui

from .documentation_window import DocumentationBuilderWindow


class DocumentationBuilder(omni.ext.IExt):
    """
    Extension to create a documentation window for any extension which has
    markdown pages defined in its extension.toml.
    Such pages can have inline omni.ui code blocks which are executed.

    For example,
    [documentation]
    pages = ["docs/Overview.md"]
    menu = "Help/API/omni.kit.documentation.builder"
    title = "Omni UI Documentation Builder"
    """

    def __init__(self):
        super().__init__()
        # cache docs and window by extensions id
        self._doc_info: typing.Dict[str, _DocInfo] = {}
        self._editor_menu = None
        self._extension_enabled_hook = None
        self._extension_disabled_hook = None

    def on_startup(self, ext_id: str):
        self._editor_menu = omni.kit.ui.get_editor_menu()
        manager = omni.kit.app.get_app().get_extension_manager()
        # add menus for any extensions with docs that are already enabled
        all_exts = manager.get_extensions()
        for ext in all_exts:
            if ext.get("enabled", False):
                self._add_docs_menu(ext.get("id"))

        # hook into extension enable/disable events if extension specifies "documentation" config key
        hooks = manager.get_hooks()
        self._extension_enabled_hook = hooks.create_extension_state_change_hook(
            self.on_after_extension_enabled,
            omni.ext.ExtensionStateChangeType.AFTER_EXTENSION_ENABLE,
            ext_dict_path="documentation",
            hook_name="omni.kit.documentation.builder",
        )
        self._extension_disabled_hook = hooks.create_extension_state_change_hook(
            self.on_before_extension_disabled,
            omni.ext.ExtensionStateChangeType.BEFORE_EXTENSION_DISABLE,
            ext_dict_path="documentation",
            hook_name="omni.kit.documentation.builder",
        )

    def on_shutdown(self):
        self._extension_enabled_hook = None
        self._extension_disabled_hook = None

        for ext_id in list(self._doc_info.keys()):
            self._destroy_docs_window(ext_id)

        self._editor_menu = None

    def on_after_extension_enabled(self, ext_id: str, *_):
        self._add_docs_menu(ext_id)

    def on_before_extension_disabled(self, ext_id: str, *_):
        self._destroy_docs_window(ext_id)

    def _add_docs_menu(self, ext_id: str):
        doc_info = _get_doc_info(ext_id)
        if doc_info and doc_info.menu_path and self._editor_menu:
            doc_info.menu = self._editor_menu.add_item(doc_info.menu_path, partial(self._show_docs_window, ext_id))
            self._doc_info[ext_id] = doc_info

    def _show_docs_window(self, ext_id: str, *_):
        window = self._doc_info[ext_id].window
        if not window:
            window = _create_docs_window(ext_id)
        self._doc_info[ext_id].window = window
        window.visible = True

    def _destroy_docs_window(self, ext_id: str):
        if ext_id in self._doc_info:
            doc_info = self._doc_info.pop(ext_id)
            if doc_info.window:
                doc_info.window.destroy()


def create_docs_window(ext_name: str):
    """
    Creates a DocumentationBuilderWindow for ext_name
    Returns the window, or None if no such extension is enabled and has documentation
    """
    # fetch by name
    ext_versions = omni.kit.app.get_app().get_extension_manager().fetch_extension_versions(ext_name)
    # use latest version that is enabled
    for ver in ext_versions:
        if ver.get("enabled", False):
            return _create_docs_window(ver.get("id", ""))
    return None


def _create_docs_window(ext_id: str):
    """
    Creates a DocumentationBuilderWindow for ext_id
    Returns the window, or None if ext_id has no documentation
    """
    doc_info = _get_doc_info(ext_id)
    if doc_info:
        window = DocumentationBuilderWindow(doc_info.title, filenames=doc_info.pages)
        return window
    return None


class _DocInfo:
    def __init__(self):
        self.pages = []
        self.menu_path = ""
        self.title = ""
        self.window = None
        self.menu = None


def _get_doc_info(ext_id: str) -> _DocInfo:
    """
    Returns extension documentation metadata if it exists.
    Returns None if exension has no documentation pages.
    """
    ext_info = omni.kit.app.get_app().get_extension_manager().get_extension_dict(ext_id)
    if ext_info:
        ext_path = ext_info.get("path", "")
        doc_dict = ext_info.get("documentation", {})
        doc_info = _DocInfo()
        doc_info.pages = doc_dict.get("pages", [])
        doc_info.menu_path = doc_dict.get("menu", "Help/API/" + ext_id)
        doc_info.title = doc_dict.get("title", ext_id)
        filenames = []
        for page in doc_info.pages:
            filenames.append(page if os.path.isabs(page) else os.path.join(ext_path, page))
        if filenames:
            doc_info.pages = filenames
            return doc_info
    return None
