import asyncio
import sys
import os
import carb.tokens

import omni.ui as ui
import omni.kit.ui
import omni.kit.app
import carb.settings
import omni.kit.menu.utils
import carb.input

import omni.kit.stage_templates as stage_templates
from omni.kit.window.title import get_main_window_title
from typing import Optional, Tuple


class ApplicationSetup:
    def __init__(self):
        self._settings = carb.settings.get_settings_interface()

        # storage for the lazy menu items
        self._lazy_menus = []

        # storage for the layout menu items
        self._layout_menu = None

        # storage for workflows
        self._workflows = []

    async def _load_layout(self, layout_file: str, keep_windows_open=False):
        try:
            from omni.kit.quicklayout import QuickLayout

            # few frames delay to avoid the conflict with the layout of omni.kit.mainwindow
            for i in range(3):
                await omni.kit.app.get_app().next_update_async()

            QuickLayout.load_file(layout_file, keep_windows_open)
            # load layout file again to make sure layout correct

            for i in range(3):
                await omni.kit.app.get_app().next_update_async()

            QuickLayout.load_file(layout_file, keep_windows_open)

        except Exception as exc:
            pass

    def force_viewport_settings(self):
        display_options = self._settings.get("/persistent/app/viewport/displayOptions")
        # Note: flags are from omni/kit/ViewportTypes.h
        kShowFlagAxis = 1 << 1
        kShowFlagGrid = 1 << 6
        kShowFlagSelectionOutline = 1 << 7
        kShowFlagLight = 1 << 8
        display_options = (
            display_options
            | (kShowFlagAxis)
            | (kShowFlagGrid)
            | (kShowFlagSelectionOutline)
            | (kShowFlagLight)
        )
        self._settings.set("/persistent/app/viewport/displayOptions", display_options)
        # Make sure these are in sync from changes above
        self._settings.set("/app/viewport/show/lights", True)
        self._settings.set("/app/viewport/grid/enabled", True)
        self._settings.set("/app/viewport/outline/enabled", True)

    def setup_application_title(self):
        # Adjust the Window Title to show the Create Version
        window_title = get_main_window_title()

        app_version = self._settings.get("/app/version")
        if not app_version:
            app_version = open(
                carb.tokens.get_tokens_interface().resolve("${app}/../VERSION")
            ).read()

        if app_version:
            if "+" in app_version:
                app_version, _ = app_version.split("+")

            # for RC version we remove some details
            production = self._settings.get("/privacy/externalBuild")
            if production:
                if "-" in app_version:
                    app_version, _ = app_version.split("-")

                if self._settings.get("/app/beta"):
                    window_title.set_app_version(app_version + " Beta")
            else:
                window_title.set_app_version(app_version)

    def load_default_layout(self, keep_windows_open=False):
        default_layout = self._settings.get("/app/layout/default")
        default_layout = carb.tokens.get_tokens_interface().resolve(default_layout)
        self.__setup_window_task = asyncio.ensure_future(
            self._load_layout(default_layout, keep_windows_open=keep_windows_open)
        )

    def __load_extension_and_show_window(self, menu_path, ext_id, window_name):
        # Remove "lazy" menu here because it could not handle menu states while window closed
        # Extension should create menu item again and handle menu states
        editor_menu = omni.kit.ui.get_editor_menu()
        editor_menu.remove_item(menu_path)

        manager = omni.kit.app.get_app().get_extension_manager()

        # if the extension is not enabled, enable it and show the window
        if not manager.is_extension_enabled(ext_id):
            manager.set_extension_enabled_immediate(ext_id, True)
            ui.Workspace.show_window(window_name, True)
        else:
            # the extensions was enabled already check for the window and toggle it
            window: ui.WindowHandle = ui.Workspace.get_window(window_name)
            if window:
                ui.Workspace.show_window(window_name, not window.visible)
            else:
                ui.Workspace.show_window(window_name, True)

    def create_trigger_menu(self):
        """ Create a menu item for each extension that has a [[trigger]] entry with a menu in it's config."""
        editor_menu = omni.kit.ui.get_editor_menu()
        manager = omni.kit.app.get_app().get_extension_manager()
        include_list = self._settings.get("/exts/omni.app.setup/lazy_menu/include_list") or []
        exclude_list = self._settings.get("/exts/omni.app.setup/lazy_menu/exclude_list") or []
        for ext in manager.get_extensions():
            # if the extensions is already enabled skip it, it already has a menu item if it need
            if ext["enabled"]:
                continue

            ext_id = ext["id"]

            info = manager.get_extension_dict(ext_id)
            for trigger in info.get("trigger", []):
                menu = trigger.get("menu", {})
                # enable to discover and debug the auto menu creation
                # print(f'menu: {menu} {menu.__class__}')
                menu_name = menu.get("name", None)
                # If the menu path not in white list or in black list, ignore
                if include_list and menu_name not in include_list:
                    continue
                if exclude_list and menu_name in exclude_list:
                    continue

                window_name = menu.get("window", None)
                if window_name is None or menu_name is None:
                    carb.log_error(f"Invalid [[trigger]] entry for extension: {ext_id}")
                    continue

                # get priority or default to 0
                priority = menu.get("priority", 0)
                self._lazy_menus.append(
                    editor_menu.add_item(
                        menu_name,
                        lambda window, value, menu_path=menu_name, ext_id=ext_id, window_name=window_name: self.__load_extension_and_show_window(
                            menu_path, ext_id, window_name
                        ),
                        toggle=True,
                        value=False,
                        priority=priority,
                    )
                )

    # TODO: this should be validated and those settings clean up
    def set_defaults(self):
        """ this is trying to setup some defaults for extensions to avoid warning """
        self._settings.set_default("/persistent/app/omniverse/bookmarks", {})
        self._settings.set_default("/persistent/app/stage/timeCodeRange", [0, 100])

        self._settings.set_default(
            "/persistent/audio/context/closeAudioPlayerOnStop", False
        )

        self._settings.set_default(
            "/persistent/app/primCreation/PrimCreationWithDefaultXformOps", True
        )
        self._settings.set_default(
            "/persistent/app/primCreation/DefaultXformOpType",
            "Scale, Rotate, Translate",
        )
        self._settings.set_default(
            "/persistent/app/primCreation/DefaultRotationOrder", "ZYX"
        )
        self._settings.set_default(
            "/persistent/app/primCreation/DefaultXformOpPrecision", "Double"
        )

        # omni.kit.property.tagging
        self._settings.set_default(
            "/persistent/exts/omni.kit.property.tagging/showAdvancedTagView", False
        )
        self._settings.set_default(
            "/persistent/exts/omni.kit.property.tagging/showHiddenTags", False
        )
        self._settings.set_default(
            "/persistent/exts/omni.kit.property.tagging/modifyHiddenTags", False
        )

        # adjust couple of viewport settings
        self._settings.set("/app/viewport/boundingBoxes/enabled", True)

    def build_layout_menu(self):
        """Build the layout menu based on registered Layouts and main views"""
        editor_menu = omni.kit.ui.get_editor_menu()

        self._layout_menu_items = []
        self._current_layout_priority = 20
        self._current_key_index = 1

        def add_layout_menu_entry(name, parameter, workflow=None):
            import inspect

            menu_path = f"Layout/{name}"
            menu = editor_menu.add_item(
                menu_path, None, False, self._current_layout_priority
            )
            self._current_layout_priority = self._current_layout_priority + 1
            hotkey = (carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, eval(f"carb.input.KeyboardInput.KEY_{self._current_key_index}")) if self._current_key_index <= 9 else None
            self._current_key_index += 1

            if inspect.isfunction(parameter):
                menu_action = omni.kit.menu.utils.add_action_to_menu(
                    menu_path,
                    lambda *_: parameter(),
                    name,
                    hotkey,
                )
            else:
                menu_action = omni.kit.menu.utils.add_action_to_menu(
                    menu_path,
                    lambda *_: asyncio.ensure_future(self._activate_layout(parameter, workflow)),
                    name,
                    hotkey,
                )

            self._layout_menu_items.append((menu, menu_action))

        # Default layout menu items
        add_layout_menu_entry(
            "Reset Layout", lambda: self.load_default_layout()
        )

        add_layout_menu_entry(
            "Viewport Only",
            "${omni.app.setup}/layouts/viewport_only.json",
        )

        # Layout menu items for workflows
        workflows = self._settings.get("/exts/omni.app.setup/workflow")
        active_now = self._settings.get("/exts/omni.app.setup/activeWorkflowAtStartup")
        self._workflows = []
        if workflows:
            for workflow in workflows:
                name = workflow.get("name", None)
                workflow_name = workflow.get("workflow", None)
                layout_name = workflow.get("layout", None)
                if name is None:
                    carb.log_error(f"Invalid [[app.workflow]] entry for extension: {workflow}")
                    continue

                self._workflows.append(workflow)
                add_layout_menu_entry(
                    " ".join([n.capitalize() for n in name.split(" ")]),
                    layout_name,
                    workflow_name,
                )

                if active_now:
                    self.active_workflow(name)

    async def _activate_layout(self, layout: str, workflow: str):
        if workflow:
            ext_manager = omni.kit.app.get_app_interface().get_extension_manager()
            ext_manager.set_extension_enabled_immediate(workflow, True)
            if not layout:
                layout = self._settings.get(f"/exts/{workflow}/default_layout")

        if layout:
            await omni.kit.app.get_app().next_update_async()
            layout_path = carb.tokens.get_tokens_interface().resolve(layout)
            if layout_path:
                await self._load_layout(layout_path)

    def active_workflow(self, name: str) -> None:
        for workflow in self._workflows:
            if workflow.get("name") == name:
                asyncio.ensure_future(self._activate_layout(workflow.get("layout"), workflow.get("workflow")))
                return
        else:
            carb.log_error(f"Cannot find workflow '{name}'!")

    def __config_menu_layout(self, title, info: dict, default_type="Menu"):
        from omni.kit.menu.utils import MenuLayout, LayoutSourceSearch

        def __split_title(title: str) -> Tuple[str, Optional[str], bool]:
            (name, source) = title.split("=") if "=" in title else (title, None)
            if name.startswith("-"):
                remove = True
                name = name[1:]
            else:
                remove = False
            return (name, source, remove)

        (name, source, remove) = __split_title(title)
        items = info.get("items", None) if info else None
        type = info.get("type", default_type) if info else default_type
        children = []
        if items:
            for item_title in items:
                (item_name, item_source, _) = __split_title(item_title)
                item_info = info.get(item_name, None)
                child = self.__config_menu_layout(item_title, item_info, default_type="Item")
                if child:
                    children.append(child)

        source_search = LayoutSourceSearch.EVERYWHERE if source else LayoutSourceSearch.LOCAL_ONLY
        if name == "":
            return MenuLayout.Seperator()
        elif type == "Menu":
            return MenuLayout.Menu(name, items=children, source=source, source_search=source_search, remove=remove)
        elif type == "SubMenu":
            return MenuLayout.SubMenu(name, items=children, source=source, source_search=source_search, remove=remove)
        elif type == "Item":
            return MenuLayout.Item(name, source=source, source_search=source_search, remove=remove)
        elif type == "Group":
            return MenuLayout.Group(name, items=children, source_search=source_search, source=source)
        elif type == "Sort":
            exclude_items = info.get("exclude_items", []) if info else []
            sort_submenus = info.get("sort_submenus", False) if info else False
            return MenuLayout.Sort(exclude_items=exclude_items, sort_submenus=sort_submenus)
        else:
            return None

    def setup_menu_layout(self):
        menu_layouts = self._settings.get("/exts/omni.app.setup/menu_layout")
        if not menu_layouts:
            return

        self._layout_menu = []
        for (menu_name, menu_info) in menu_layouts.items():
            menu = self.__config_menu_layout(menu_name, menu_info)
            if menu:
                self._layout_menu.append(menu)

        omni.kit.menu.utils.add_layout(self._layout_menu)

    def setup_menu_order(self):
        menu_order_presets = self._settings.get("/exts/omni.app.setup/menu_order")
        if menu_order_presets:
            menu_self = omni.kit.menu.utils.get_instance()
            _, menu_order, _ = menu_self.get_menu_data()

            menu_order.update(menu_order_presets)
            menu_self.rebuild_menus()
