# 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.
#

__all__ = [
    'frame_viewport_prims',
    'frame_viewport_selection',
    'get_viewport_from_window_name',
    'get_active_viewport',
    'get_active_viewport_window'
    'get_active_viewport_and_window',
    'get_viewport_window_camera_path',
    'get_viewport_window_camera_string',
    'get_active_viewport_camera_path',
    'get_active_viewport_camera_string',
    'get_num_viewports',
    'capture_viewport_to_file',
    'capture_viewport_to_buffer',
    'post_viewport_message',
    'toggle_global_visibility',
    'create_drop_helper',
    'disable_selection',
    'get_ground_plane_info'
]


import asyncio
import carb

from pxr import Gf, Sdf

from typing import Callable, List, Optional, Tuple


_g_is_viewport_next = None


def _is_viewport_next(self):
    global _g_is_viewport_next
    if _g_is_viewport_next is None:
        vp_window_name = carb.settings.get_settings().get('/exts/omni.kit.viewport.window/startup/windowName')
        _g_is_viewport_next = vp_window_name and (vp_window_name == 'Viewport')
    return _g_is_viewport_next


def _get_default_viewport_window_name(window_name: str = None):
    if window_name:
        return window_name
    return carb.settings.get_settings().get('/exts/omni.kit.viewport.window/startup/windowName') or 'Viewport'


def get_viewport_from_window_name(window_name: str = None):
    '''Return the first Viewport that is held inside a specific Window name.'''
    window_name = _get_default_viewport_window_name(window_name)
    try:
        from omni.kit.viewport.window import get_viewport_window_instances
        # Get every ViewportWindow, regardless of UsdContext it is attached to
        for window in get_viewport_window_instances(None):
            if window.title == window_name:
                return window.viewport_api
    except ImportError:
        pass

    try:
        import omni.kit.viewport_legacy as vp_legacy
        vp_iface = vp_legacy.get_viewport_interface()
        viewport_handle = vp_iface.get_instance(window_name)
        if viewport_handle:
            vp_window = vp_iface.get_viewport_window(viewport_handle)
            if vp_window:
                from .legacy_viewport_api import LegacyViewportAPI
                return LegacyViewportAPI(vp_iface.get_viewport_window_name(viewport_handle))
    except ImportError:
        pass

    return None


def get_active_viewport_and_window(usd_context_name: str = '', wrap_legacy: bool = True, window_name: str = None):
    '''Return the active Viewport for a given UsdContext and the name of the Window it is inside of.'''
    default_window_name = _get_default_viewport_window_name(window_name)
    try:
        from omni.kit.viewport.window import get_viewport_window_instances, ViewportWindow
        # If no windowname provided, see if the ViewportWindow already knows what is active
        if window_name is None:
            active_window = ViewportWindow.active_window
            if active_window:
                # Have an active Window, need to make sure UsdContext name matches (or passed None to avoid the match)
                viewport_api = active_window.viewport_api
                if (usd_context_name is None) or (usd_context_name == viewport_api.usd_context_name):
                    return (viewport_api, active_window)

        active_window = None
        # Get all ViewportWindows attached the UsdContext with this name
        for window in get_viewport_window_instances(usd_context_name):
            # If matching by name, check that first ignoring whether focused or not (multiple Windows cannot have same name)
            window_title = window.title
            if window_name and window_name != window_title:
                continue

            # If this Window is focused, then return it
            if window.focused:
                active_window = window
                break
            # Save the first encountered Window as he fallback 'default' Window
            if window_title == default_window_name:
                active_window = window
            elif active_window is None:
                active_window = window
        if active_window:
            return (active_window.viewport_api, active_window)
    except ImportError:
        pass

    try:
        import omni.kit.viewport_legacy as vp_legacy
        vp_iface = vp_legacy.get_viewport_interface()
        instance_list = vp_iface.get_instance_list()
        if instance_list:
            first_context_match = None
            for viewport_handle in vp_iface.get_instance_list():
                vp_window = vp_iface.get_viewport_window(viewport_handle)
                if not vp_window:
                    continue
                # If matching by name, check that first ignoring whether focused or not (multiple Windows cannot have same name)
                window_title = vp_iface.get_viewport_window_name(viewport_handle)
                if window_name and window_name != vp_iface.get_viewport_window_name(viewport_handle):
                    continue

                # Filter by UsdContext name, where None means any UsdContext
                if (usd_context_name is not None) and (usd_context_name != vp_window.get_usd_context_name()):
                    continue
                if vp_window.is_focused():
                    first_context_match = viewport_handle
                    break
                elif window_title == default_window_name:
                    first_context_match = viewport_handle
                elif first_context_match is None:
                    first_context_match = viewport_handle

            # If there was a match on UsdContext name (but not focused), return the first one
            if first_context_match is not None:
                vp_window = vp_iface.get_viewport_window(first_context_match)
                if vp_window:
                    from .legacy_viewport_api import LegacyViewportAPI
                    window_name = vp_iface.get_viewport_window_name(first_context_match)
                    viewport_api = LegacyViewportAPI(vp_iface.get_viewport_window_name(first_context_match))
                    if wrap_legacy:
                        from .legacy_viewport_window import LegacyViewportWindow
                        vp_window = LegacyViewportWindow(window_name, viewport_api)
                    return (viewport_api, vp_window)
    except ImportError:
        pass

    return (None, None)


def get_active_viewport_window(window_name: str = None, wrap_legacy: bool = True, usd_context_name: str = ''):
    '''Return the active Viewport for a given UsdContext.'''
    return get_active_viewport_and_window(usd_context_name, wrap_legacy, window_name)[1]


def get_active_viewport(usd_context_name: str = ''):
    '''Return the active Viewport for a given UsdContext.'''
    return get_active_viewport_and_window(usd_context_name, False)[0]


def get_viewport_window_camera_path(window_name: str = None) -> Sdf.Path:
    '''Return a Sdf.Path to the camera used by the active Viewport in a named Window.'''
    viewport_api = get_viewport_from_window_name(window_name)
    return viewport_api.camera_path if viewport_api else None


def get_viewport_window_camera_string(window_name: str = None) -> str:
    '''Return a string path to the camera used by the active Viewport in a named Window.'''
    viewport_api = get_viewport_from_window_name(window_name)
    return viewport_api.camera_path.pathString if viewport_api else None


def get_active_viewport_camera_path(usd_context_name: str = '') -> Sdf.Path:
    '''Return a Sdf.Path to the camera used by the active Viewport for a specific UsdContext.'''
    viewport_api = get_active_viewport(usd_context_name)
    return viewport_api.camera_path if viewport_api else None


def get_active_viewport_camera_string(usd_context_name: str = '') -> str:
    '''Return a string path to the camera used by the active Viewport for a specific UsdContext.'''
    viewport_api = get_active_viewport(usd_context_name)
    return viewport_api.camera_path.pathString if viewport_api else None


def get_available_aovs_for_viewport(viewport_api):
    if hasattr(viewport_api, 'legacy_window'):
        viewport_handle = viewport_api.frame_info.get('viewport_handle')
        asyncio.ensure_future(viewport_api.usd_context.next_frame_async(viewport_handle))
        return viewport_api.legacy_window.get_aov_list()

    carb.log_error('Available AOVs not implemented')
    return []


def add_aov_to_viewport(viewport_api, aov_name: str):
    if hasattr(viewport_api, 'legacy_window'):
        return viewport_api.legacy_window.add_aov(aov_name)

    from pxr import Usd, UsdRender
    from omni.usd import editor

    stage  = viewport_api.stage
    render_product_path = viewport_api.render_product_path
    with Usd.EditContext(stage, stage.GetSessionLayer()):
        render_prod_prim = stage.GetPrimAtPath(render_product_path)
        if not render_prod_prim:
            raise RuntimeError(f'Invalid renderProduct "{render_product_path}"')
        render_var_prim_path = Sdf.Path(f'/Render/Vars/{aov_name}')
        render_var_prim = stage.GetPrimAtPath(render_var_prim_path)
        if not render_var_prim:
            render_var_prim = stage.DefinePrim(render_var_prim_path)
        if not render_var_prim:
            raise RuntimeError(f'Cannot create renderVar "{render_var_prim_path}"')
        render_var_prim.CreateAttribute("sourceName", Sdf.ValueTypeNames.String).Set(aov_name)
        render_prod_var_rel = render_prod_prim.GetRelationship('orderedVars')
        if not render_prod_var_rel:
            render_prod_prim.CreateRelationship('orderedVars')
        if not render_prod_var_rel:
            raise RuntimeError(f'cannot set orderedVars relationship for renderProduct "{render_product_path}"')
        render_prod_var_rel.AddTarget(render_var_prim_path)

        editor.set_hide_in_stage_window(render_var_prim, True)
        editor.set_no_delete(render_var_prim, True)

    return True


def post_viewport_message(viewport_api_or_window, message: str, message_id: str = None):
    if hasattr(viewport_api_or_window, 'legacy_window'):
        viewport_api_or_window.legacy_window.post_toast(message)
        return
    if hasattr(viewport_api_or_window, '_post_toast_message'):
        viewport_api_or_window._post_toast_message(message, message_id)
        return
    try:
        from omni.kit.viewport.window import get_viewport_window_instances
        for window in get_viewport_window_instances(viewport_api_or_window.usd_context_name):
            if window.viewport_api.id == viewport_api_or_window.id:
                window._post_toast_message(message, message_id)
                return
    except (ImportError, AttributeError):
        pass


class _CaptureHelper:
    def __init__(self, legacy_window, is_hdr: bool, file_path: str = None, render_product_path: str = None, on_capture_fn: Callable = None, format_desc: dict = None):
        import omni.renderer_capture
        self.__future = asyncio.Future()
        self.__renderer = omni.renderer_capture.acquire_renderer_capture_interface()
        if render_product_path:
            self.__future.set_result(True)
            self.__renderer.capture_next_frame_using_render_product(viewport_handle=legacy_window.get_id(), filepath=file_path, render_product=render_product_path)
            return

        self.__is_hdr = is_hdr
        self.__legacy_window = legacy_window
        self.__file_path = file_path
        self.__on_capture_fn = on_capture_fn
        self.__format_desc = format_desc
        event_stream = legacy_window.get_ui_draw_event_stream()
        self.__capture_sub = event_stream.create_subscription_to_pop(self.capture_function, name='omni.kit.viewport.utility.capture_viewport')

    def capture_function(self, *args):
        self.__capture_sub = None
        legacy_window, self.__legacy_window = self.__legacy_window, None
        renderer, self.__renderer = self.__renderer, None
        if self.__is_hdr:
            viewport_rp_resource = legacy_window.get_drawable_hdr_resource()
        else:
            viewport_rp_resource = legacy_window.get_drawable_ldr_resource()

        if self.__on_capture_fn:
            def _interecept_capture(*args, **kwargs):
                try:
                    self.__on_capture_fn(*args, **kwargs)
                finally:
                    if not self.__future.done():
                        self.__future.set_result(True)
            renderer.capture_next_frame_rp_resource_callback(_interecept_capture, resource=viewport_rp_resource)
        elif self.__file_path:
            if not self.__future.done():
                self.__future.set_result(True)

            if self.__format_desc:
                if hasattr(renderer, 'capture_next_frame_rp_resource_to_file'):
                    renderer.capture_next_frame_rp_resource_to_file(filepath=self.__file_path,
                                                                    resource=viewport_rp_resource,
                                                                    format_desc=self.__format_desc)
                    return
                carb.log_error('Format description provided to capture, but not honored')

            renderer.capture_next_frame_rp_resource(filepath=self.__file_path, resource=viewport_rp_resource)

    async def wait_for_result(self, completion_frames: int = 2):
        await self.__future
        import omni.kit.app
        app = omni.kit.app.get_app()
        while completion_frames:
            await app.next_update_async()
            completion_frames = completion_frames - 1
        return self.__future.result()


def capture_viewport_to_buffer(viewport_api, on_capture_fn: Callable, is_hdr: bool = False):
    '''Capture the provided viewport and send it to a callback.'''
    if hasattr(viewport_api, 'legacy_window'):
        return _CaptureHelper(viewport_api.legacy_window, is_hdr=is_hdr, on_capture_fn=on_capture_fn)

    from omni.kit.widget.viewport.capture import ByteCapture
    return viewport_api.schedule_capture(ByteCapture(on_capture_fn, aov_name='HdrColor' if is_hdr else 'LdrColor'))


def capture_viewport_to_file(viewport_api, file_path: str = None, is_hdr: bool = False, render_product_path: str = None, format_desc: dict = None):
    '''Capture the provided viewport to a file.'''
    file_path = str(file_path)
    if hasattr(viewport_api, 'legacy_window'):
        return _CaptureHelper(viewport_api.legacy_window, is_hdr=is_hdr, file_path=file_path,
                              render_product_path=render_product_path, format_desc=format_desc)

    from omni.kit.widget.viewport.capture import MultiAOVFileCapture
    class HdrCaptureHelper(MultiAOVFileCapture):
        def __init__(self, file_path: str, is_hdr: bool, format_desc: dict = None):
            super().__init__(['HdrColor' if is_hdr else 'LdrColor'], [file_path])
            # Setup RenderProduct for Hdr

        def __del__(self):
            # Setdown RenderProduct for Hdr
            pass

        def capture_aov(self, file_path, aov):
            if render_product_path:
                self.save_product_to_file(file_path, render_product_path)
            else:
                self.save_aov_to_file(file_path, aov, format_desc=format_desc)

    return viewport_api.schedule_capture(HdrCaptureHelper(file_path, is_hdr))

def get_num_viewports(usd_context_name: str = None):
    num_viewports = 0
    try:
        from omni.kit.viewport.window import get_viewport_window_instances
        num_viewports += sum(1 for _ in get_viewport_window_instances(usd_context_name))
    except ImportError:
        pass

    try:
        import omni.kit.viewport_legacy as vp_legacy
        vp_iface = vp_legacy.get_viewport_interface()
        for viewport_handle in vp_iface.get_instance_list():
            if usd_context_name and (usd_context_name != vp_iface.get_viewport_window(viewport_handle).get_usd_context_name()):
                continue
            num_viewports += 1
    except ImportError:
        pass

    return num_viewports


def create_viewport_window(name: str = None, usd_context_name: str = '', width: int = 1280, height: int = 720, position_x: int = 0, position_y: int = 0, camera_path: Sdf.Path = None, **kwargs):
    window = None

    try:
        from omni.kit.viewport.window import get_viewport_window_instances, ViewportWindow
        if name is None:
            name = f"Viewport {get_num_viewports()}"

        window = ViewportWindow(name, usd_context_name, width=width, height=height, **kwargs)

        return window
    except ImportError:
        pass

    if window is None:
        try:
            import omni.kit.viewport_legacy as vp_legacy
            from .legacy_viewport_window import LegacyViewportWindow
            vp_iface = vp_legacy.get_viewport_interface()
            vp_handle = vp_iface.create_instance()
            assigned_name = vp_iface.get_viewport_window_name(vp_handle)
            if name and (name != assigned_name):
                carb.log_warn('omni.kit.viewport_legacy does not support explicit Window names, using assigned name "{assigned_name}"')
            window = LegacyViewportWindow(assigned_name, **kwargs)
            window.width = width
            window.height = height
        except ImportError:
            pass

    if window:
        window.setPosition(position_x, position_y)
        if camera_path:
            window.viewport_api.camera_path = camera_path

    return window


class ViewportPrimReferencePoint:
    BOUND_BOX_CENTER = 0
    BOUND_BOX_LEFT = 1
    BOUND_BOX_RIGHT = 2
    BOUND_BOX_TOP = 3
    BOUND_BOX_BOTTOM = 4


def get_ui_position_for_prim(viewport_window, prim_path: str, alignment: ViewportPrimReferencePoint = ViewportPrimReferencePoint.BOUND_BOX_CENTER, force_legacy_api: bool = False):
    if isinstance(viewport_window, str):
        window_name = str(viewport_window)
        viewport_window = get_active_viewport_window(window_name=window_name)
        if viewport_window is None:
            carb.log_error('No ViewportWindow found with name "{window_name}"')
            return (0, 0), False

    # XXX: omni.ui constants needed
    import omni.ui
    dpi = omni.ui.Workspace.get_dpi_scale()
    if dpi <= 0.0:
        dpi = 1

    # XXX: kit default dock splitter size (4)
    dock_splitter_size = 4 * dpi
    tab_bar_height = 0
    if viewport_window.dock_tab_bar_visible or not (viewport_window.flags & omni.ui.WINDOW_FLAGS_NO_TITLE_BAR):
        tab_bar_height = 22 * dpi

    # Force legacy Viewport code path if requested
    if force_legacy_api and hasattr(viewport_window, 'legacy_window'):
        import omni.ui
        import omni.kit.viewport_legacy as vp_legacy

        legacy_window = viewport_window.legacy_window
        alignment = {
            ViewportPrimReferencePoint.BOUND_BOX_LEFT: vp_legacy.ViewportPrimReferencePoint.BOUND_BOX_LEFT,
            ViewportPrimReferencePoint.BOUND_BOX_RIGHT: vp_legacy.ViewportPrimReferencePoint.BOUND_BOX_RIGHT,
            ViewportPrimReferencePoint.BOUND_BOX_TOP: vp_legacy.ViewportPrimReferencePoint.BOUND_BOX_TOP,
            ViewportPrimReferencePoint.BOUND_BOX_BOTTOM: vp_legacy.ViewportPrimReferencePoint.BOUND_BOX_BOTTOM,
        }.get(alignment, vp_legacy.ViewportPrimReferencePoint.BOUND_BOX_CENTER)

        success, x, y, z = legacy_window.get_prim_clipping_pos(str(prim_path), alignment)
        if not success:
            return (0, 0), False

        # Changing xy to match the window coord system.
        x = (x + 1.0) / 2.0   # x to [0, 1]
        y = 1 - (y + 1.0) / 2.0  # y to [0, 1] and reverse its direction.

        min_x, min_y, max_x, max_y = legacy_window.get_viewport_rect()
        prim_window_pos_x = (max_x - min_x) * x + min_x - dock_splitter_size
        prim_window_pos_y = (max_y - min_y) * y + min_y - tab_bar_height - dock_splitter_size
    else:
        from pxr import UsdGeom, Gf
        viewport_api = viewport_window.viewport_api
        usd_context = viewport_window.viewport_api.usd_context
        stage = usd_context.get_stage()
        usd_prim = stage.GetPrimAtPath(prim_path) if stage else False
        if usd_prim:
            xformable_prim = UsdGeom.Xformable(usd_prim)
        else:
            xformable_prim = None
        if (not stage) or (not xformable_prim):
            return (0, 0), False

        # Get bounding box from prim
        aabb_min, aabb_max = usd_context.compute_path_world_bounding_box(str(prim_path))
        gf_range = Gf.Range3d(Gf.Vec3d(aabb_min[0], aabb_min[1], aabb_min[2]), Gf.Vec3d(aabb_max[0], aabb_max[1], aabb_max[2]))
        if gf_range.IsEmpty():
            # May be empty (UsdGeom.Xform for example), so build a scene-scaled constant box
            world_units = UsdGeom.GetStageMetersPerUnit(stage)
            if Gf.IsClose(world_units, 0.0, 1e-6):
                world_units = 0.01

            # XXX: compute_path_world_transform is identity in this case
            if False:
                world_xform = Gf.Matrix4d(*usd_context.compute_path_world_transform(str(prim_path)))
            else:
                import omni.timeline
                time = omni.timeline.get_timeline_interface().get_current_time() * stage.GetTimeCodesPerSecond()
                world_xform = xformable_prim.ComputeLocalToWorldTransform(time)

            ref_position = world_xform.ExtractTranslation()
            extent = Gf.Vec3d(0.2 / world_units) # 20cm by default
            gf_range.SetMin(ref_position - extent)
            gf_range.SetMax(ref_position + extent)

        # Computes the extent in clipping pos
        mvp = viewport_api.world_to_ndc
        min_x, min_y, min_z = 2.0, 2.0, 2.0
        max_x, max_y, max_z = -2.0, -2.0, -2.0
        for i in range(8):
            corner = gf_range.GetCorner(i)
            pos = mvp.Transform(corner)
            min_x = min(min_x, pos[0])
            min_y = min(min_y, pos[1])
            min_z = min(min_z, pos[2])
            max_x = max(max_x, pos[0])
            max_y = max(max_y, pos[1])
            max_z = max(max_z, pos[2])

        min_point = Gf.Vec3d(min_x, min_y, min_z)
        max_point = Gf.Vec3d(max_x, max_y, max_z)
        mid_point = (min_point + max_point) / 2

        # Map to reference point in screen space
        if alignment == ViewportPrimReferencePoint.BOUND_BOX_LEFT:
            ndc_pos = mid_point - Gf.Vec3d((max_x - min_x) / 2, 0, 0)
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_RIGHT:
            ndc_pos = mid_point + Gf.Vec3d((max_x - min_x) / 2, 0, 0)
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_TOP:
            ndc_pos = mid_point + Gf.Vec3d(0, (max_y - min_y) / 2, 0)
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_BOTTOM:
            ndc_pos = mid_point - Gf.Vec3d(0, (max_y - min_y) / 2, 0)
        else:
            ndc_pos = mid_point

        # Make sure its not clipped
        if (ndc_pos[2] < 0) or (ndc_pos[0] < -1) or (ndc_pos[0] > 1) or (ndc_pos[1] < -1) or (ndc_pos[1] > 1):
            return (0, 0), False

        '''
        XXX: Simpler world calculation
        world_pos = gf_range.GetMidpoint()

        dir_sel = (0, 1)
        up_axis = UsdGeom.GetStageUpAxis(stage)
        if up_axis == UsdGeom.Tokens.z:
            dir_sel = (0, 2)
        if up_axis == UsdGeom.Tokens.x:
            dir_sel = (2, 1)

        if alignment == ViewportPrimReferencePoint.BOUND_BOX_LEFT:
            world_pos[dir_sel[0]] = gf_range.GetMin()[dir_sel[0]]
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_RIGHT:
            world_pos[dir_sel[0]] = gf_range.GetMax()[dir_sel[0]]
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_TOP:
            world_pos[dir_sel[1]] = gf_range.GetMax()[dir_sel[1]]
        elif alignment == ViewportPrimReferencePoint.BOUND_BOX_BOTTOM:
            world_pos[dir_sel[1]] = gf_range.GetMin()[dir_sel[1]]

        ndc_pos = viewport_api.world_to_ndc.Transform(world_pos)
        '''

        x = (ndc_pos[0] + 1.0) / 2.0   # x to [0, 1]
        y = 1 - (ndc_pos[1] + 1.0) / 2.0  # y to [0, 1] and reverse its direction.

        frame = viewport_window.frame
        prim_window_pos_x = dpi * (frame.screen_position_x + x * frame.computed_width) - dock_splitter_size
        prim_window_pos_y = dpi * (frame.screen_position_y + y * frame.computed_height) - tab_bar_height - dock_splitter_size

    return (prim_window_pos_x, prim_window_pos_y), True


def frame_viewport_prims(viewport_api=None, prims: List[str] = None):
    if not prims:
        return False
    return __frame_viewport_objects(viewport_api, prims=prims, force_legacy_api=False)


def frame_viewport_selection(viewport_api=None, force_legacy_api: bool = False):
    return __frame_viewport_objects(viewport_api, prims=None, force_legacy_api=force_legacy_api)


def __frame_viewport_objects(viewport_api=None, prims: Optional[List[str]] = None, force_legacy_api: bool = False):
    if not viewport_api:
        viewport_api = get_active_viewport()
    if not viewport_api:
        return False

    if force_legacy_api and hasattr(viewport_api, "legacy_window"):
        viewport_api.legacy_window.focus_on_selected()
        return

    # This is new CODE
    if prims is None:
        # Get current selection
        prims = viewport_api.usd_context.get_selection().get_selected_prim_paths()
        # Pass None to underlying command to signal "frame all" if selection is empty
        prims = prims if prims else None

    stage = viewport_api.stage
    cam_path = viewport_api.camera_path
    if not stage or not cam_path:
        return False
    cam_prim = stage.GetPrimAtPath(cam_path)
    if not cam_prim:
        return False

    import omni.kit.undo
    import omni.kit.commands
    from pxr import UsdGeom

    look_through = None
    # Loop over all targets (should really be only one) and see if we can get a valid UsdGeom.Imageable
    for target in cam_prim.GetRelationship('omni:kit:viewport:lookThrough:target').GetForwardedTargets():
        target_prim = stage.GetPrimAtPath(target)
        if not target_prim:
            continue
        if UsdGeom.Imageable(target_prim):
            look_through = target_prim
            break

    try:
        omni.kit.undo.begin_group()
        resolution = viewport_api.resolution
        omni.kit.commands.execute(
            'FramePrimsCommand',
            prim_to_move=cam_path if not look_through else look_through.GetPath(),
            prims_to_frame=prims,
            time_code=viewport_api.time,
            usd_context_name=viewport_api.usd_context_name,
            aspect_ratio=resolution[0] / resolution[1]
        )
    finally:
        omni.kit.undo.end_group()

    return True


def toggle_global_visibility(force_legacy_api: bool = False):
    # Forward to omni.kit.viewport.actions and propogate an errors so caller should update code
    carb.log_warn("omni.kit.viewport.utility.toggle_global_visibility is deprecated, use omni.kit.viewport.actions.toggle_global_visibility")
    import omni.kit.actions.core
    action_registry = omni.kit.actions.core.get_action_registry()
    action_registry.get_action("omni.kit.viewport.actions", "toggle_global_visibility").execute()


async def next_viewport_frame_async(viewport, n_frames: int = 0):
    """
    Waits until frames have been delivered to the Viewport.

    Args:
        viewport: the Viewport to wait for a frame on.
        n_frames: the number of rendered frames to wait for.
    """
    import omni.kit.app
    app = omni.kit.app.get_app()
    # Make sure at least one frame has been delivered
    viewport_handle = viewport.frame_info.get('viewport_handle')
    while viewport_handle is None:
        await app.next_update_async()
        viewport_handle = viewport.frame_info.get('viewport_handle')
    # Now wait for any additional frames
    usd_context = viewport.usd_context
    while n_frames:
        await usd_context.next_frame_async(viewport_handle)
        n_frames = n_frames - 1


def create_drop_helper(*args, **kwargs):
    try:
        from omni.kit.viewport.window.dragdrop.legacy import create_drop_helper
        return create_drop_helper(*args, **kwargs)
    except (ImportError, ModuleNotFoundError):
        try:
            import omni.kit.viewport_legacy as vp_legacy
            return vp_legacy.get_viewport_interface().create_drop_helper(*args, **kwargs)
        except (ImportError, ModuleNotFoundError):
            pass


class _DisableViewportWindowLayer:
    def __init__(self, viewport_or_window, layers_and_categories):
        self.__layers = []

        # Find the omni.ui.Window for this Viewport to disable selection manipulator
        viewport_window = None
        from omni.kit.viewport.window import ViewportWindow, get_viewport_window_instances
        if not isinstance(viewport_or_window, ViewportWindow):
            viewport_id = viewport_or_window.id
            for window_instance in get_viewport_window_instances(viewport_or_window.usd_context_name):
                viewport_window_viewport = window_instance.viewport_api
                if viewport_window_viewport.id == viewport_id:
                    viewport_window = window_instance
                    break
        else:
            viewport_window = viewport_or_window

        if viewport_window:
            for layer, category in layers_and_categories:
                found_layer = viewport_window._find_viewport_layer(layer, category)
                if found_layer:
                    self.__layers.append((found_layer, found_layer.visible))
                    found_layer.visible = False

    def __del__(self):
        for layer, visible in self.__layers:
            layer.visible = visible
        self.__layers = tuple()


def disable_selection(viewport_or_window, disable_click: bool = True):
    '''Disable selection rect and possible the single click selection on a Viewport or ViewportWindow.
    Returns an object that resets selection when it goes out of scope.'''

    # First check if the Viewport is a legacy Viewport
    from .legacy_viewport_window import LegacyViewportWindow
    legacy_window = None
    if hasattr(viewport_or_window, 'legacy_window'):
        legacy_window = viewport_or_window.legacy_window
    elif isinstance(viewport_or_window, LegacyViewportWindow):
        legacy_window = viewport_or_window.legacy_window

    if legacy_window:
        class LegacySelectionState:
            def __init__(self, viewport_window, disable_click):
                self.__viewport_window = viewport_window
                self.__restore_picking = viewport_window.is_enabled_picking()
                if disable_click:
                    self.__viewport_window.set_enabled_picking(False)
                self.__viewport_window.disable_selection_rect(True)

            def __del__(self):
                self.__viewport_window.set_enabled_picking(self.__restore_picking)

        return LegacySelectionState(legacy_window, disable_click)

    disable_items = [('Selection', 'manipulator')]
    if disable_click:
        disable_items.append(('ObjectClick', 'manipulator'))

    return _DisableViewportWindowLayer(viewport_or_window, disable_items)


def disable_context_menu(viewport_or_window=None):
    '''Disable context menu on a Viewport or ViewportWindow.
    Returns an object that resets context menu visibility when it goes out of scope.'''

    if viewport_or_window:
        # Check if the Viewport is a legacy Viewport, s can only operate on all Viewportwindows in that case
        from .legacy_viewport_window import LegacyViewportWindow
        if hasattr(viewport_or_window, 'legacy_window') or isinstance(viewport_or_window, LegacyViewportWindow):
            carb.log_warn("Cannot disable context menu on an individual Viewport, disabling it for all")
            viewport_or_window = None

    if viewport_or_window is None:
        class _DisableAllContextMenus:
            def __init__(self):
                # When setting is initially unset, then context menu is enabled
                self.__settings = carb.settings.get_settings()
                enabled = self.__settings.get('/exts/omni.kit.window.viewport/showContextMenu')
                self.__settings.set('/exts/omni.kit.window.viewport/showContextMenu', False)
                self.__restore = enabled if (enabled is not None) else True

            def __del__(self):
                self.__settings.set('/exts/omni.kit.window.viewport/showContextMenu', self.__restore)

        return _DisableAllContextMenus()

    return _DisableViewportWindowLayer(viewport_or_window, [('ContextMenu', 'manipulator')])


def get_ground_plane_info(viewport, ortho_special: bool = True) -> Tuple[Gf.Vec3d, List[str]]:
    """For a given Viewport returns a tuple representing the ground plane.

    @param viewport: ViewportAPI to use for ground plane info
    @param ortho_special: bool Use an alternate ground plane for orthographic cameras that are looking down a single axis
    @return: A tuple of type (normal: Gf.Vec3d, planes: List[str]), where planes is a list of axis' occupied (x, y, z)
    """

    stage = viewport.stage
    cam_path = viewport.camera_path

    from pxr import UsdGeom
    up_axis = UsdGeom.GetStageUpAxis(stage) if stage else UsdGeom.Tokens.y
    if up_axis == UsdGeom.Tokens.y:
        normal = Gf.Vec3d.YAxis()
        planes = ['x', 'z']
    elif up_axis == UsdGeom.Tokens.z:
        normal = Gf.Vec3d.ZAxis()
        planes = ['x', 'y']
    else:
        normal = Gf.Vec3d.XAxis()
        planes = ['y', 'z']

    if ortho_special:
        cam_prim = stage.GetPrimAtPath(cam_path) if stage else None
        if cam_prim:
            usd_camera = UsdGeom.Camera(cam_prim)
            if usd_camera and usd_camera.GetProjectionAttr().Get(viewport.time) == UsdGeom.Tokens.orthographic:
                orthoEpsilon = 0.0001
                viewNormal = viewport.transform.TransformDir(Gf.Vec3d(0, 0, -1))
                if Gf.IsClose(abs(viewNormal[1]), 1.0, orthoEpsilon):
                    normal = Gf.Vec3d.YAxis()
                    planes = ['x', 'z']
                elif Gf.IsClose(abs(viewNormal[2]), 1.0, orthoEpsilon):
                    normal = Gf.Vec3d.ZAxis()
                    planes = ['x', 'y']
                elif Gf.IsClose(abs(viewNormal[0]), 1.0, orthoEpsilon):
                    normal = Gf.Vec3d.XAxis()
                    planes = ['y', 'z']

    return normal, planes
