import os
import asyncio
import carb.input
import omni.ext
import omni.usd
import omni.ui as ui
import omni.kit.app
import omni.kit.menu.utils
import omni.kit.notification_manager as nm
import omni.kit.usd.layers as layers
from functools import partial
from omni.kit.menu.utils import MenuItemDescription
from .selection import Selection
from .selection_window import SelectionSetWindow
from .edit_actions import register_actions, deregister_actions
from pxr import Sdf, Usd, Kind, Tf

_extension_instance = None
_extension_path = None

PERSISTENT_SETTINGS_PREFIX = "/persistent"
KEEP_TRANSFORM_FOR_REPARENTING = "/persistent/app/stage/movePrimInPlace"


class EditMenuExtension(omni.ext.IExt):
    def __init__(self):
        super().__init__()
        omni.kit.menu.utils.set_default_menu_proirity("Edit", -9)
        self._select_recent_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._selection_set_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._selection_kind_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._select_recent = []
        self._selection_set = []

    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, EditMenuExtension, lambda: _extension_instance)

        self._select_recent_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._selection_set_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._selection_kind_menu_list = [ MenuItemDescription(name="None", enabled=False) ]
        self._select_recent = []
        self._selection_set = []
        self._create_selection_window = None
        self._edit_menu_list = None
        self._build_edit_menu()
#        self._stage_event_sub = omni.usd.get_context().get_stage_event_stream().create_subscription_to_pop(self._on_stage_event, name="omni.usd.menu.edit")

        self._context_menus = []
        manager = omni.kit.app.get_app().get_extension_manager()
        self._cm_hooks = manager.subscribe_to_extension_enable(
            lambda _: self._register_context_menu(),
            lambda _: self._unregister_context_menu(),
            ext_name="omni.kit.context_menu",
            hook_name="menu edit listener",
        )

    def on_shutdown(self):
        global _extension_instance
        global _template_list

        deregister_actions(self._ext_name)

        self._unregister_context_menu()
        self._cm_hooks = None

        _extension_instance = None
        _template_list = None
#        self._stage_event_sub = None
        self._select_recent_menu_list = []
        self._selection_set_menu_list = []
        self._selection_kind_menu_list = []
        self._select_recent = []
        self._selection_set = []

        if self._create_selection_window:
            self._create_selection_window.shutdown()
            del self._create_selection_window

        omni.kit.menu.utils.remove_menu_items(self._edit_menu_list, "Edit")

    def post_notification(message: str, info: bool = False, duration: int = 3):
        if info:
            type = nm.NotificationStatus.INFO
        else:
            type = nm.NotificationStatus.WARNING

        nm.post_notification(message, status=type, duration=duration)

    def _usd_kinds(self):
        return ["model", "assembly", "group", "component", "subcomponent"]

    def _usd_kinds_display(self):
        return self._usd_kinds()[1:]

    def _plugin_kinds(self):
        all_kinds = set(Kind.Registry.GetAllKinds())
        return all_kinds - set(self._usd_kinds())

    def _create_selection_set(self):
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Create Selection Set as no prims are selected")
            return

        if self._create_selection_window:
            self._create_selection_window.show()
        else:
            self._create_selection_window = SelectionSetWindow(self._add_to_selection_set)

    def _build_recent_menu(self):
        if self._select_recent:
            self._select_recent_menu_list.clear()

            self._select_recent.sort(reverse=True, key=lambda x: x.time)
            del self._select_recent[5:]

            for recent in self._select_recent:
                index = self._select_recent.index(recent)
                self._select_recent_menu_list.append(
                    MenuItemDescription(
                        name=recent.description, onclick_action=("omni.kit.menu.edit", "select_recent", index)
                    )
                )

            omni.kit.menu.utils.rebuild_menus()

    def _on_select_recent(self, index):
        async def select_func():
            await omni.kit.app.get_app().next_update_async()
            do_recent = self._select_recent[index]
            omni.kit.commands.execute("SelectListCommand", selection=do_recent.selection)
            do_recent.touch()
            self._build_recent_menu()

        asyncio.ensure_future(select_func())

    def _add_to_selection_set(self, description):
        new_selection = omni.usd.get_context().get_selection().get_selected_prim_paths()
        for selection in self._selection_set:
            if selection.description == description:
                EditMenuExtension.post_notification(f"Selection set \"{description}\" already exists")
                return

        self._selection_set.append(Selection(description, new_selection))
        self._build_selection_set_menu()

    def _add_to_recent(self, description, selection):
        found = False
        for prev_selection in self._select_recent:
            if prev_selection.description == description and prev_selection.selection == selection:
                prev_selection.touch()
                found = True
                break

        if not found:
            self._select_recent.append(Selection(description, selection))

        self._build_recent_menu()

    def _on_select_selection_set(self, index):
        async def select_func():
            await omni.kit.app.get_app().next_update_async()
            do_select = self._selection_set[index]
            omni.kit.commands.execute("SelectListCommand", selection=do_select.selection)

        asyncio.ensure_future(select_func())

    def _build_selection_set_menu(self):
        self._selection_set_menu_list.clear()

        self._selection_set.sort(key=lambda x: x.description)

        for recent in self._selection_set:
            index=self._selection_set.index(recent)
            self._selection_set_menu_list.append(MenuItemDescription(name=recent.description, onclick_action=("omni.kit.menu.edit", "select_selection_set", index)))

        omni.kit.menu.utils.rebuild_menus()

    def _on_select_by_kind(self, kind):
        async def select_func():
            await omni.kit.app.get_app().next_update_async()
            omni.kit.commands.execute("SelectKindCommand", kind=kind)

        asyncio.ensure_future(select_func())

    def _build_selection_kind_menu(self):
        self._selection_kind_menu_list.clear()

        for kind in self._usd_kinds_display():
            self._selection_kind_menu_list.append(MenuItemDescription(name=str(kind).capitalize(), onclick_action=("omni.kit.menu.edit", "select_by_kind", kind)))

        for kind in self._plugin_kinds():
            self._selection_kind_menu_list.append(MenuItemDescription(name=str(kind).capitalize(), onclick_action=("omni.kit.menu.edit", "select_by_kind", kind)))

    def _build_edit_menu(self):
        self._build_recent_menu()
        self._build_selection_kind_menu()
        # setup menus
        self._edit_select_menu = [
            MenuItemDescription(
                name="Select Recent",
                sub_menu=self._select_recent_menu_list
            ),
            MenuItemDescription(
                name="Select All",
                onclick_action=("omni.kit.selection", "all"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.A),
            ),
            MenuItemDescription(
                name="Select None",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.selection", "none"),
                hotkey=(0, carb.input.KeyboardInput.ESCAPE),
            ),
            MenuItemDescription(
                name="Select Invert",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "selection_invert"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.I),
            ),
            MenuItemDescription(
                name="Select Parent",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "selection_parent"),
                hotkey=(0, carb.input.KeyboardInput.UP),
            ),
            MenuItemDescription(
                name="Select Leaf",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "selection_leaf"),
            ),
            MenuItemDescription(
                name="Select Hierarchy",
                onclick_action=("omni.kit.menu.edit", "selection_hierarchy"),
                # enable_fn=EditMenuExtension.prim_selected,
            ),
            MenuItemDescription(
                name="Select Similar",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "selection_similar"),
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Create Selection Set",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "create_selection_set"),
            ),
            MenuItemDescription(
                name="Select Set",
                sub_menu=self._selection_set_menu_list
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Select by Kind",
                sub_menu=self._selection_kind_menu_list
            )]
        self._edit_menu_list = [
            MenuItemDescription(
                name="Undo",
                glyph="none.svg",
                onclick_action=("omni.kit.commands", "undo"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.Z),
            ),
            MenuItemDescription(
                name="Redo",
                glyph="none.svg",
                onclick_action=("omni.kit.commands", "redo"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.Y),
            ),
            MenuItemDescription(
                name="Repeat",
                glyph="none.svg",
                onclick_action=("omni.kit.commands", "repeat"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.R),
            ),
            MenuItemDescription(
                name="Select",
                glyph="none.svg",
                sub_menu=self._edit_select_menu
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Instance",
                glyph="none.svg",
                # enable_fn=[EditMenuExtension.prim_selected, EditMenuExtension.can_be_instanced],
                onclick_action=("omni.kit.menu.edit", "instance_prim"),
                hotkey=(
                    carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL | carb.input.KEYBOARD_MODIFIER_FLAG_SHIFT,
                    carb.input.KeyboardInput.I,
                ),
            ),
            MenuItemDescription(
                name="Duplicate",
                glyph="copy.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "duplicate_prim"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.D),
            ),
            MenuItemDescription(
                name="Duplicate - All Layers",
                glyph="copy.svg",
                # enable_fn=[EditMenuExtension.prim_selected, EditMenuExtension.is_not_in_live_session],
                onclick_action=("omni.kit.menu.edit", "duplicate_prim_and_layers"),
            ),
            MenuItemDescription(
                name="Duplicate - Collapsed",
                glyph="copy.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "duplicate_prim_and_combine_layers"),
            ),
            MenuItemDescription(
                name="Parent",
                glyph="copy.svg",
                # enable_fn=[EditMenuExtension.can_prims_parent, EditMenuExtension.is_not_in_live_session],
                onclick_action=("omni.kit.menu.edit", "parent_prims"),
                hotkey=(0, carb.input.KeyboardInput.P),
            ),
            MenuItemDescription(
                name="Unparent",
                glyph="copy.svg",
                # enable_fn=[EditMenuExtension.can_prims_unparent, EditMenuExtension.is_not_in_live_session],
                onclick_action=("omni.kit.menu.edit", "unparent_prims"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_SHIFT, carb.input.KeyboardInput.P),
            ),
            MenuItemDescription(
                name="Group",
                glyph="none.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "create_xform_to_group"),
                hotkey=(carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL, carb.input.KeyboardInput.G),
            ),
            MenuItemDescription(
                name="Ungroup",
                glyph="none.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "ungroup_prims"),
            ),
            MenuItemDescription(
                name="Toggle Visibility",
                glyph="none.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "toggle_visibillity"),
                hotkey=(0, carb.input.KeyboardInput.H),
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Delete",
                glyph="trash.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "delete_prim"),
                hotkey=(0, carb.input.KeyboardInput.DEL),
            ),
            MenuItemDescription(
                name="Delete - All Layers",
                glyph="trash.svg",
                # enable_fn=[EditMenuExtension.prim_selected, EditMenuExtension.is_not_in_live_session],
                onclick_action=("omni.kit.menu.edit", "delete_prim_all_layers"),
            ),
            MenuItemDescription(
                name="Rename",
                glyph="pencil.svg",
                # enable_fn=[EditMenuExtension.prim_selected, EditMenuExtension.is_one_prim_selected, EditMenuExtension.can_delete],
                onclick_action=("omni.kit.menu.edit", "menu_rename_prim_dialog"),
                hotkey=(0, carb.input.KeyboardInput.F2),
                appear_after="Delete - All Layers",
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Hide Unselected",
                glyph="none.svg",
                onclick_action=("omni.kit.selection", "HideUnselected"),
            ),
            MenuItemDescription(
                name="Unhide All",
                glyph="none.svg",
                onclick_action=("omni.kit.selection", "UnhideAllPrims"),
            ),
            MenuItemDescription(),
            MenuItemDescription(
                name="Focus",
                glyph="none.svg",
                # enable_fn=EditMenuExtension.prim_selected,
                onclick_action=("omni.kit.menu.edit", "focus_prim"),
                hotkey=(0, carb.input.KeyboardInput.F),
            ),
            MenuItemDescription(
                name="Toggle Visualization Mode",
                glyph="none.svg",
                onclick_action=("omni.kit.viewport.actions", "toggle_global_visibility"),
            ),
            MenuItemDescription(
                name="Capture Screenshot",
                glyph="none.svg",
                onclick_action=("omni.kit.menu.edit", "capture_screenshot"),
                hotkey=(0, carb.input.KeyboardInput.F10),
            ),
        ]

        omni.kit.menu.utils.add_menu_items(self._edit_menu_list, "Edit", -9)

#    def _on_stage_event(self, event):
#        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED):
#            async def refresh_menu_items():
#                omni.kit.menu.utils.refresh_menu_items("Edit")

#            asyncio.ensure_future(refresh_menu_items())

    def _register_context_menu(self):
        import omni.kit.context_menu
        select_context_menu_dict = {
            "name": {
                "Select": [
                    {"name": {"Select Recent": self._select_recent_menu_list}},
                    {"name": "Select All", "onclick_action": ("omni.kit.selection", "all")},
                    {"name": "Select None", "onclick_action": ("omni.kit.selection", "none")},
                    {"name": "Select Invert", "onclick_action": ("omni.kit.menu.edit", "selection_invert")},
                    {"name": "Select Parent", "onclick_action": ("omni.kit.menu.edit", "selection_parent")},
                    {"name": "Select Leaf", "onclick_action": ("omni.kit.menu.edit", "selection_leaf")},
                    {"name": "Select Hierarchy", "onclick_action": ("omni.kit.menu.edit", "selection_hierarchy")},
                    {"name": "Select Similar", "onclick_action": ("omni.kit.menu.edit", "selection_similar")},
                    {"name": ""},
                    {"name": {"Select by Kind": self._selection_kind_menu_list}},
                ]
            },
            "glyph": "none.svg",
            "appear_after": ""
        }
        self._context_menus.append(
            omni.kit.context_menu.add_menu(select_context_menu_dict, "MENU", "omni.kit.window.viewport")
        )

    def _unregister_context_menu(self):
        for menu in self._context_menus:
            menu.release()
        self._context_menus.clear()

    def prim_selected():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        return bool(paths)

    def is_in_live_session():
        usd_context = omni.usd.get_context()
        live_syncing = layers.get_layers(usd_context).get_live_syncing()
        return live_syncing.is_stage_in_live_session()

    def is_not_in_live_session():
        return not EditMenuExtension.is_stage_in_live_session()

    def can_delete():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        stage = omni.usd.get_context().get_stage()
        if not stage:
            return False
        for path in paths:
            prim = stage.GetPrimAtPath(path)
            if not prim or not prim.IsValid():
                return False
            no_delete = prim.GetMetadata("no_delete")
            if no_delete is not None and no_delete is True:
                return False
        return True

    def is_one_prim_selected():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        return len(paths) == 1

    def can_prims_parent():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if paths:
            return bool(len(paths)>1)
        return False

    def can_prims_unparent():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if paths:
            return True
        return False

    def can_be_instanced():
        allowed_types_for_instancing = ["Xform"]
        stage = omni.usd.get_context().get_stage()
        if not stage:
            return False
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        for path in paths:
            prim = stage.GetPrimAtPath(path)
            if not prim or prim.GetTypeName() not in allowed_types_for_instancing:
                return False
        return True

    def instance_prim():
        usd_context = omni.usd.get_context()
        paths = usd_context.get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot instance prim as no prims are selected")
            return

        with omni.kit.usd.layers.active_authoring_layer_context(usd_context):
            paths = usd_context.get_selection().get_selected_prim_paths()
            omni.kit.commands.execute("CreateInstances", paths_from=paths)

    # When combine_layers is True, it means to duplicate with flattening references.
    def duplicate_prim(duplicate_layers, combine_layers):
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot duplicate prim as no prims are selected")
            return

        if EditMenuExtension.is_in_live_session() and duplicate_layers and not combine_layers:
            warn = "Destructive duplication is not supported in live-syncing."
            carb.log_warn(warn)
            EditMenuExtension.post_notification(warn)
            return

        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        usd_context = omni.usd.get_context()
        with layers.active_authoring_layer_context(usd_context):
            omni.kit.commands.execute(
                "CopyPrims", paths_from=paths, duplicate_layers=duplicate_layers, combine_layers=combine_layers,
                flatten_references=combine_layers
            )

    def parent_prims():
        stage = omni.usd.get_context().get_stage()
        if not stage:
            EditMenuExtension.post_notification("Cannot parent prim as no stage is loaded")
            return

        if not EditMenuExtension.can_prims_parent():
            EditMenuExtension.post_notification("Cannot parent prim as two or more prims are not selected")
            return

        if EditMenuExtension.is_in_live_session():
            EditMenuExtension.post_notification("Parent is not supported in live-syncing.")
            return

        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if paths:
            def prim_renamed(old_prim_name: Sdf.Path, new_prim_name: Sdf.Path):
                try:
                    index = paths.index(old_prim_name.pathString)
                    paths[index] = new_prim_name.pathString
                except Exception as exc:
                    carb.log_warn(f"prim_renamed error {exc}")

            parent_path = paths.pop()
            prim = stage.GetPrimAtPath(parent_path)
            if prim and prim.IsInstanceable():
                EditMenuExtension.post_notification(f"{parent_path} cannot be parent as its Instanceable")
                return

            settings = carb.settings.get_settings()
            keep_transform = settings.get(KEEP_TRANSFORM_FOR_REPARENTING)
            if keep_transform is None:
                keep_transform = True

            omni.kit.commands.execute(
                "ParentPrims",
                parent_path=parent_path,
                child_paths=paths,
                on_move_fn=prim_renamed,
                keep_world_transform=keep_transform
            )
            omni.usd.get_context().get_selection().set_selected_prim_paths(paths, True)

    def unparent_prims():
        stage = omni.usd.get_context().get_stage()
        if not stage:
            EditMenuExtension.post_notification("Cannot unparent prim as no stage is loaded")
            return

        if not EditMenuExtension.can_prims_unparent():
            EditMenuExtension.post_notification("Cannot unparent prim as no prims are not selected")
            return

        if EditMenuExtension.is_in_live_session():
            EditMenuExtension.post_notification("Unparent is not supported in live-syncing.")
            return

        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if paths:
            settings = carb.settings.get_settings()
            keep_transform = settings.get(KEEP_TRANSFORM_FOR_REPARENTING)
            if keep_transform is None:
                keep_transform = True

            omni.kit.commands.execute("UnparentPrims", paths=paths, keep_world_transform=keep_transform)

    def create_xform_to_group():
        usd_context = omni.usd.get_context()
        paths = usd_context.get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Group as no prims are selected")
            return

        stage = usd_context.get_stage()
        if not stage:
            EditMenuExtension.post_notification("Cannot Group prim as no stage is loaded")
            return

        if stage.HasDefaultPrim() and stage.GetDefaultPrim().GetPath().pathString in paths:
            EditMenuExtension.post_notification("Cannot Group default prim")
            return

        with omni.kit.usd.layers.active_authoring_layer_context(usd_context):
            omni.kit.commands.execute("GroupPrims", prim_paths=paths, destructive=False)
        return True

    def ungroup_prims():
        usd_context = omni.usd.get_context()
        paths = usd_context.get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Ungroup as no prims are selected")
            return

        stage = usd_context.get_stage()
        if not stage:
            EditMenuExtension.post_notification("Cannot Ungroup prim as no stage is loaded")
            return

        with omni.kit.usd.layers.active_authoring_layer_context(usd_context):
            omni.kit.commands.execute("UngroupPrims", prim_paths=paths, destructive=False)
        return True

    def delete_prim(destructive):
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Delete as no prims are selected")
            return

        if destructive and EditMenuExtension.is_in_live_session():
            EditMenuExtension.post_notification("Delete prims is not supported in live-syncing.")
            return

        omni.kit.commands.execute("DeletePrims", paths=paths, destructive=destructive)

    def focus_prim():
        from omni.kit.viewport.utility import frame_viewport_selection
        frame_viewport_selection(force_legacy_api=True)

    def toggle_visibillity():
        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Toggle Visibillity as no prims are selected")
            return

        omni.kit.commands.execute("ToggleVisibilitySelectedPrims", selected_paths=paths)

    def toggle_global_visibility():
        carb.log_error("omni.kit.menu.edit.toggle_global_visibility is deprecated, use omni.kit.viewport.actions.toggle_global_visibility")
        action_registry = omni.kit.actions.core.get_action_registry()
        exc_action = action_registry.get_action("omni.kit.viewport.actions", "toggle_global_visibility")
        if exc_action:
            return exc_action.execute()

        carb.log_error("omni.kit.viewport.actions must be enabled")

    def capture_screenshot(on_complete_fn: callable=None):
        from datetime import datetime

        now = datetime.now()
        date_time = now.strftime("%Y-%m-%d %H.%M.%S")
        capture_filename = "capture.%s.png" % (date_time)

        # Check if the user specified the screenshots folder.
        settings = carb.settings.get_settings()
        screenshot_path = settings.get(PERSISTENT_SETTINGS_PREFIX + "/app/captureFrame/path")
        if screenshot_path:
            if os.path.isdir(screenshot_path):
                capture_filename = os.path.join(screenshot_path, capture_filename)
            else:
                carb.log_error(f"Can't save screenshot to {str(screenshot_path)} because it is not a directory")


        async def capture_frame(capture_filename: str, on_complete_fn: callable):
            import omni.kit.app

            carb.log_warn("Capturing %s" % (capture_filename))

            # wait for next frame so menu has closed
            await omni.kit.app.get_app().next_update_async()

            capture_viewport = settings.get(PERSISTENT_SETTINGS_PREFIX + "/app/captureFrame/viewport")
            success = False
            try:
                if not capture_viewport:
                    import omni.renderer_capture

                    renderer_capture = omni.renderer_capture.acquire_renderer_capture_interface()
                    renderer_capture.capture_next_frame_swapchain(capture_filename)
                    success = True
                else:
                    from omni.kit.viewport.utility import get_active_viewport, capture_viewport_to_file

                    viewport = get_active_viewport()
                    if viewport:
                        capture_viewport_to_file(viewport, file_path=capture_filename, is_hdr=False)
                        success = True
            except ImportError:
                pass

            if not success:
                carb.log_error("Failed to capture " + ("viewport" if capture_viewport else "") + "screenshot")

            if on_complete_fn:
                # wait for screenshot to complete
                await omni.kit.app.get_app().next_update_async()
                await omni.kit.app.get_app().next_update_async()
                on_complete_fn(success, capture_filename)

        # wait for next frame so menu has closed
        asyncio.ensure_future(capture_frame(capture_filename, on_complete_fn))

    @staticmethod
    def _rename_viewport_active_camera(old_prim_name: Sdf.Path, new_prim_name: Sdf.Path):
        try:
            from omni.kit.viewport.utility import get_active_viewport
            viewport = get_active_viewport()
            if viewport and viewport.camera_path == old_prim_name:
                viewport.camera_path = new_prim_name or "/OmniverseKit_Persp"
        except ImportError:
            pass

    @staticmethod
    def rename_prim(stage, prim, window, field_widget):
        if window:
            window.visible = False
        if field_widget:
            def select_new_prim(old_prim_name: Sdf.Path, new_prim_name: Sdf.Path):
                omni.usd.get_context().get_selection().set_selected_prim_paths([new_prim_name.pathString], True)

            if prim.GetPath().name != field_widget.model.get_value_as_string():
                old_prim_name = prim.GetPath()
                new_prim_name = prim.GetPath().GetParentPath()
                new_prim_name = new_prim_name.AppendChild(Tf.MakeValidIdentifier(field_widget.model.get_value_as_string()))
                new_prim_name = omni.usd.get_stage_next_free_path(stage, new_prim_name.pathString, False)
                on_move_fn = select_new_prim

                if Sdf.Path.IsValidPathString(new_prim_name):
                    EditMenuExtension._rename_viewport_active_camera(old_prim_name, new_prim_name)
                    move_dict = {old_prim_name: new_prim_name}
                    omni.kit.commands.execute("MovePrims", paths_to_move=move_dict, on_move_fn=on_move_fn, destructive=False)
                else:
                    EditMenuExtension.post_notification(f"Cannot rename {old_prim_name} to {new_prim_name} as its not a valid USD path")

    def menu_rename_prim_dialog(self):
        import omni.usd

        paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
        if not paths:
            EditMenuExtension.post_notification("Cannot Rename prim as no prims are selected")
            return

        if len(paths) != 1:
            EditMenuExtension.post_notification("Cannot Rename more then one prim")
            return

        stage = omni.usd.get_context().get_stage()
        if not stage or not paths:
            return

        # try to use omni.kit.widget.stage rename...
        try:
            import omni.kit.window.stage

            window = ui.Workspace.get_window("Stage")
            stage_widget = window.get_widget()
            items = stage_widget.get_treeview().selection
            if len(items) == 1:
                return stage_widget.get_delegate().get_name_column_delegate().rename_item(items[0])

        except Exception as exc:
            carb.log_warn(f"stage window rename failed due to {exc}")

        # do dialog rename
        prim = stage.GetPrimAtPath(paths[0])
        window = ui.Window(
            "Rename " + prim.GetPath().name + "###context_menu_rename",
            width=200,
            height=0,
            flags=ui.WINDOW_FLAGS_NO_RESIZE | ui.WINDOW_FLAGS_NO_SCROLLBAR | ui.WINDOW_FLAGS_MODAL,
        )
        with window.frame:
            with ui.VStack(
                height=0,
                spacing=5,
                name="top_level_stack",
                style={"VStack::top_level_stack": {"margin": 5}, "Button": {"margin": 0}},
            ):
                async def focus(field):
                    await omni.kit.app.get_app().next_update_async()
                    field.focus_keyboard()

                new_name_widget = ui.StringField()
                new_name_widget.model.set_value(prim.GetPath().name)
                new_name_widget.focus_keyboard()

                ui.Spacer(width=5, height=5)
                with ui.HStack(spacing=5):
                    ui.Button(
                        "Ok",
                        clicked_fn=partial(EditMenuExtension.rename_prim, stage, prim, window, new_name_widget),
                    )
                    ui.Button(
                        "Cancel", clicked_fn=partial(EditMenuExtension.rename_prim, stage, prim, window, None)
                    )

                    editing_started = False

                    def on_begin():
                        nonlocal editing_started
                        editing_started = True

                    def on_end():
                        nonlocal editing_started
                        editing_started = False

                    def window_pressed_key(key_index, key_flags, key_down):
                        nonlocal editing_started
                        key_mod = key_flags & ~ui.Widget.FLAG_WANT_CAPTURE_KEYBOARD
                        if carb.input.KeyboardInput(key_index) in [carb.input.KeyboardInput.ENTER, carb.input.KeyboardInput.NUMPAD_ENTER] and key_mod == 0 and key_down and editing_started:
                            on_end()
                            EditMenuExtension.rename_prim(stage, prim, window, new_name_widget)

                    new_name_widget.model.add_begin_edit_fn(lambda m: on_begin())
                    new_name_widget.model.add_end_edit_fn(lambda m: on_end())
                    window.set_key_pressed_fn(window_pressed_key)
                    asyncio.ensure_future(focus(new_name_widget))


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
