import asyncio
import functools
import carb
import carb.settings
import omni.ext
import omni.kit.app
from functools import partial
from typing import Callable, Tuple, Union, List
from .actions import add_action_to_menu, ActionMenuSubscription
from .layout import MenuLayout
from .builder_utils import MenuItemDescription, PrebuiltItemOrder, MenuAlignment, LayoutSourceSearch
from .app_menu import MenuItemOrder, MenuState, MenuActionControl, IconMenuDelegate

_extension_instance = None

class MenuUtilsExtension(omni.ext.IExt):
    def __init__(self):
        super().__init__()
        self._menu_defs = {}
        self._menu_delegates = {}
        self._menu_order = {}
        self._master_menu = {}
        self._ready_state = MenuState.Invalid
        self._menus_to_refresh = set()
        self._refresh_task = None

        # debug stats
        self._stats = {}
        self._stats["add_menu_items"] = 0
        self._stats["remove_menu_items"] = 0
        self._stats["add_hook"] = 0
        self._stats["remove_hook"] = 0
        self._stats["add_layout"] = 0
        self._stats["remove_layout"] = 0
        self._stats["refresh_menu_items"] = 0
        self._stats["refresh_menu_items_skipped"] = 0
        self._stats["rebuild_menus"] = 0
        self._stats["rebuild_menus_skipped"] = 0

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

        LOADED_EXTENSION_COUNT_MENU = "/exts/omni.kit.menu.utils/loaded_extension_count"
        settings = carb.settings.get_settings()
        settings.set_default_int(LOADED_EXTENSION_COUNT_MENU, 0)
        loaded_extension_count = settings.get(LOADED_EXTENSION_COUNT_MENU) + 1
        settings.set(LOADED_EXTENSION_COUNT_MENU, loaded_extension_count)

        self._menu_creator = None
        self._menu_layout = MenuLayout()
        self._hooks = []
        self._stats["extension_loaded_count"] = loaded_extension_count

        ext_manager = omni.kit.app.get_app_interface().get_extension_manager()

        # Hook to extension enable/disable
        # to setup hotkey with different ways with omni.kit.hotkeys.core enabled/disabled
        hooks = ext_manager.get_hooks()
        self.__extension_enabled_hook = hooks.create_extension_state_change_hook(
            self.__on_ext_changed, omni.ext.ExtensionStateChangeType.AFTER_EXTENSION_ENABLE
        )
        self.__extension_disabled_hook = hooks.create_extension_state_change_hook(
            self.__on_ext_changed, omni.ext.ExtensionStateChangeType.AFTER_EXTENSION_DISABLE
        )

        mainwindow_loaded = next(
            (
                ext
                for ext in ext_manager.get_extensions()
                if ext["id"].startswith("omni.kit.mainwindow") and ext["enabled"]
            ),
            None,
        )

        settings = carb.settings.get_settings()
        if settings.get("/exts/omni.kit.menu.utils/forceEditorMenu") is not None:
            carb.log_error(f"omni.kit.menu.utils forceEditorMenu is no longer supported")
 
        if mainwindow_loaded:
            from .app_menu import AppMenu

            self._menu_creator = AppMenu(lambda: _extension_instance)
        else:
            carb.log_info(f"omni.kit.mainwindow is not loaded. Menus are disabled")

            # if mainwindow loads later, need to refresh menu_creator
            manager = omni.kit.app.get_app().get_extension_manager()
            self._hooks.append(
                manager.subscribe_to_extension_enable(
                    on_enable_fn=lambda _: self._mainwindow_loaded(),
                    on_disable_fn=None,
                    ext_name="omni.kit.mainwindow",
                    hook_name="omni.kit.menu.utils omni.kit.mainwindow listener",
                )
            )

        # set app started trigger. refresh_menu_items & rebuild_menus won't do anything until self._ready_state is MenuState.Created
        self._app_ready_sub = (
            omni.kit.app.get_app()
            .get_startup_event_stream()
            .create_subscription_to_pop_by_type(
                omni.kit.app.EVENT_APP_READY, self._build_menus_after_loading, name="omni.kit.menu.utils app started trigger"
            )
        )

    def __on_ext_changed(self, ext_id: str, *_):
        if ext_id.startswith("omni.kit.hotkeys.core"):
            self.rebuild_menus()

    def _mainwindow_loaded(self):
        from .app_menu import AppMenu

        self._menu_creator = AppMenu(lambda: _extension_instance)
        self.rebuild_menus()
        carb.log_info(f"omni.kit.mainwindow is now loaded. Menus are enabled")

    def _build_menus_after_loading(self, event):
        self._ready_state = MenuState.Created
        self._menu_layout.menus_created()
        self.rebuild_menus()

    def on_shutdown(self):
        global _extension_instance
        _extension_instance = None

        self.__extension_enabled_hook = None
        self.__extension_disabled_hook = None
        if self._menu_creator:
            self._menu_creator.destroy()
            del self._menu_creator
            self._menu_creator = None

        if self._menu_layout:
            self._menu_layout.destroy()
            del self._menu_layout
            self._menu_layout = None

        self._menus_to_refresh = None
        self._refresh_task = None
        self._menu_defs = None
        self._menu_delegates = None
        self._menu_order = None
        self._hooks = None

    def add_menu_items(self, menu: list, name: str, menu_index: int, rebuild_menus: bool, delegate = None) -> list:
        if not menu and delegate:
            menu = [MenuItemDescription(name="placeholder", show_fn=lambda: False)]

        if menu and not isinstance(menu[0], MenuItemDescription):
            carb.log_error(f"add_menu_items: menu {menu} is not a MenuItemDescription")
            return None

        self._stats["add_menu_items"] += 1

        if not name in self._menu_defs:
            self._menu_defs[name] = []
        self._menu_defs[name].append(menu)

        if name in self._menu_delegates:
            _delegate, count = self._menu_delegates[name]
            if delegate != _delegate:
                carb.log_warn(f"add_menu_items: menu {menu} cannot change delegate")
            self._menu_delegates[name] = (_delegate, count + 1)
        elif delegate:
            self._menu_delegates[name] = (delegate, 1)

        if menu_index != 0 and name not in self._menu_order:
            self._menu_order[name] = menu_index

        if rebuild_menus:
            self.rebuild_menus()

        return menu

    def set_default_menu_proirity(self, name: str, menu_index: int):
        if menu_index == 0:
            return

        for index in self._menu_order:
            if self._menu_order[index] == menu_index:
                return
        self._menu_order[name] = menu_index

    def remove_menu_items(self, menu: list, name: str, rebuild_menus: bool):
        self._stats["remove_menu_items"] += 1
        try:
            if name in self._menu_defs:
                self._menu_defs[name].remove(menu)

            if name in self._menu_delegates:
                delegate, count = self._menu_delegates[name]
                if count == 1:
                    del self._menu_delegates[name]
                else:
                    self._menu_delegates[name] = (delegate, count - 1)

            if rebuild_menus:
                self.rebuild_menus()
        except Exception as exc:
            carb.log_warn(f"omni.kit.menu.utils remove_menu_items \"{name}\" failed {exc}")

    def add_hook(self, callback: Callable):
        self._stats["add_hook"] += 1
        if self._menu_creator:
            self._menu_creator.add_hook(callback)

    def remove_hook(self, callback: Callable):
        self._stats["remove_hook"] += 1
        try:
            if self._menu_creator:
                self._menu_creator.remove_hook(callback)
        except Exception as exc:
            carb.log_warn(f"omni.kit.menu.utils remove_hook failed {exc}")

    def add_layout(self, layout: List[Union[MenuLayout.Menu, MenuLayout.SubMenu, MenuLayout.Item, MenuLayout.Seperator, MenuLayout.Group]]):
        self._stats["add_layout"] += 1
        self._menu_layout.add_layout(layout)
        self.rebuild_menus()

    def remove_layout(self, layout: List[Union[MenuLayout.Menu, MenuLayout.SubMenu, MenuLayout.Item, MenuLayout.Seperator, MenuLayout.Group]]):
        self._stats["remove_layout"] += 1
        self._menu_layout.remove_layout(layout)
        self.rebuild_menus()

    def get_merged_menus(self):
        return self._menu_creator.merge_menus(self._menu_defs.keys(), self._menu_defs, self._menu_order)

    def refresh_menu_items(self, name: str, immediately: bool = False):
        if not self._ready_state & MenuState.Created:
            self._stats["refresh_menu_items_skipped"] += 1
            return

        self._stats["refresh_menu_items"] += 1
        if self._menu_creator:
            self._menu_creator.refresh_menu_items(name)

    def rebuild_menus(self):
        if not self._ready_state & MenuState.Created:
            self._stats["rebuild_menus_skipped"] += 1
            return

        self._stats["rebuild_menus"] += 1

        if self._menu_creator:
            self._menu_creator.create_menu()

    def get_menu_layout(self):
        return self._menu_layout

    def get_menu_data(self):
        return self._menu_defs, self._menu_order, self._menu_delegates


def get_instance():
    return _extension_instance


def add_menu_items(menu: list, name: str, menu_index: int = 0, rebuild_menus: bool = True, delegate = None):
    """
    add a list of menus items to menu.
    menu is list of MenuItemDescription()
    name is name to appear when menu is collapsed
    menu_index is horizontal positioning
    rebuild_menus is flag to call rebuild_menus when True
    delegate ui.MenuDelegate delegate
    """
    instance = get_instance()
    if instance:
        return get_instance().add_menu_items(menu, name, menu_index, rebuild_menus, delegate)


def remove_menu_items(menu: list, name: str, rebuild_menus: bool = True):
    """
    remove  a list of menus items to menu.
    menu is list of MenuItemDescription()
    name is name to appear when menu is collapsed
    rebuild_menus is flag to call rebuild_menus when True
    """
    instance = get_instance()
    if instance:
        instance.remove_menu_items(menu, name, rebuild_menus)


def refresh_menu_items(name: str, immediately = None):
    """
    update menus enabled state
    menu is list of MenuItemDescription()
    name is name to appear when menu is collapsed
    immediately is deprecated and not used
    """
    instance = get_instance()
    if instance:
        carb.log_info(f"omni.kit.menu.utils.refresh_menu_items {name}")
        if immediately != None:
            carb.log_warn(f"refresh_menu_items immediately parameter is deprecated and not used")
        instance.refresh_menu_items(name)


def add_hook(callback: Callable):
    """
    add a menu modification callback hook
    callback is function to be called when menus are re-generated
    """
    instance = get_instance()
    if instance:
        instance.add_hook(callback)


def remove_hook(callback: Callable):
    """
    remove a menu modification callback hook
    callback is function to be called when menus are re-generated
    """
    instance = get_instance()
    if instance:
        instance.remove_hook(callback)


def rebuild_menus():
    """
    force menus to rebuild, triggering hooks
    """
    instance = get_instance()
    if instance:
        carb.log_info(f"omni.kit.menu.utils.rebuild_menus")
        instance.rebuild_menus()


def set_default_menu_proirity(name, menu_index):
    """
    set default menu priority
    """
    instance = get_instance()
    if instance:
        instance.set_default_menu_proirity(name, menu_index)


def add_layout(layout: List[Union[MenuLayout.Menu, MenuLayout.SubMenu, MenuLayout.Item, MenuLayout.Seperator, MenuLayout.Group]]):
    """
    add a menu layout.
    """
    instance = get_instance()
    if instance:
        instance.add_layout(layout)


def remove_layout(layout: List[Union[MenuLayout.Menu, MenuLayout.SubMenu, MenuLayout.Item, MenuLayout.Seperator, MenuLayout.Group]]):
    """
    remove a menu layout.
    """
    instance = get_instance()
    if instance:
        instance.remove_layout(layout)


def get_menu_layout():
    """
    get menu layouts.
    """
    instance = get_instance()
    if instance:
        return instance.get_menu_layout().get_layout()


def get_merged_menus() -> dict:
    """
    get combined menus as dictionary
    """
    instance = get_instance()
    if instance:
        return instance.get_merged_menus()

    return None


def get_debug_stats() -> dict:
    instance = get_instance()
    if instance:
        return instance._stats

    return None
