# Copyright (c) 2021-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__ = ['LegacyViewportAPI']


import carb
import omni.ui
import omni.usd
import omni.timeline
import omni.kit.app
from pxr import Gf, Sdf, Usd, UsdGeom, CameraUtil
from typing import Callable, List, Tuple, Sequence


# Wrap a legacy viewport Window in an object that poses as a new ViewportAPI

class LegacyViewportAPI:
    def __init__(self, window_name: str):
        self.__window_name = window_name
        self.__settings = carb.settings.get_settings()
        # All legacy Viewports are locked to timeline time
        self.__timeline = None
        self.__saved_resolution = None
        self.__fill_frame = None

    @property
    def camera_path(self) -> Sdf.Path:
        '''Return an Sdf.Path to the active rendering camera'''
        path_str = self.legacy_window.get_active_camera()
        return Sdf.Path(path_str) if path_str else Sdf.Path()

    @camera_path.setter
    def camera_path(self, camera_path: Sdf.Path):
        '''Set the active rendering camera from an Sdf.Path'''
        self.legacy_window.set_active_camera(str(camera_path))

    @property
    def render_product_path(self) -> str:
        '''Return a string to the UsdRender.Product used by the Viewport'''
        return self.legacy_window.get_render_product_path()

    @render_product_path.setter
    def render_product_path(self, product_path: str):
        '''Set the UsdRender.Product used by the Viewport with a string'''
        self.legacy_window.set_render_product_path(str(product_path))

    @property
    def resolution(self) -> Tuple[float]:
        '''Return a tuple of (resolution_x, resolution_y) this Viewport is rendering at.'''
        # Legacy Viewport applies scale internally, so returned value accounts for scale factor.
        return self.legacy_window.get_texture_resolution()

    @resolution.setter
    def resolution(self, value: Tuple[float]):
        '''Set the resolution to render with (resolution_x, resolution_y).'''
        # Legacy Viewport applies scale internally based on single carb setting, so just use that
        self.legacy_window.set_texture_resolution(value[0], value[1])

    @property
    def resolution_scale(self) -> float:
        '''Get the scaling factor for the Viewport's render resolution.'''
        # Legacy Viewport applies scale internally based on single carb setting, so just use that or fallback to 1
        return self.__settings.get("/app/renderer/resolution/multiplier") or 1.0

    @resolution_scale.setter
    def resolution_scale(self, value: float):
        '''Set the scaling factor for the Viewport's render resolution.'''
        value = float(value)
        if value <= 0:
            raise ValueError("Viewport resolution scale must be greater than 0.")
        self.__settings.set("/app/renderer/resolution/multiplier", value)

    @property
    def full_resolution(self) -> Tuple[float]:
        '''Return a tuple of the full (full_resolution_x, full_resolution_y) this Viewport is rendering at, not accounting for scale.'''
        # Legacy Viewport applies scale internally, so undo any scaling factor
        resolution = self.resolution
        resolution_scale = self.resolution_scale
        return (resolution[0] / resolution_scale, resolution[1] / resolution_scale)

    # Legacy methods that we also support
    def get_active_camera(self) -> Sdf.Path:
        '''Return an Sdf.Path to the active rendering camera'''
        return self.camera_path

    def set_active_camera(self, camera_path: Sdf.Path):
        '''Set the active rendering camera from an Sdf.Path'''
        self.camera_path = camera_path

    def get_render_product_path(self) -> str:
        '''Return a string to the UsdRender.Product used by the Viewport'''
        return self.render_product_path

    def set_render_product_path(self, product_path: str):
        '''Set the UsdRender.Product used by the Viewport with a string'''
        self.render_product_path = product_path

    def get_texture_resolution(self) -> Tuple[float]:
        '''Return a tuple of (resolution_x, resolution_y)'''
        return self.resolution

    def set_texture_resolution(self, value: Tuple[float]):
        '''Set the resolution to render with (resolution_x, resolution_y)'''
        self.resolution = value

    def get_texture_resolution_scale(self) -> float:
        '''Get the scaling factor for the Viewport's render resolution.'''
        return self.resolution_scale

    def set_texture_resolution_scale(self, value: float):
        '''Set the scaling factor for the Viewport's render resolution.'''
        self.resolution_scale = value

    def get_full_texture_resolution(self) -> Tuple[float]:
        '''Return a tuple of (full_resolution_x, full_resolution_y)'''
        return self.full_resolution

    @property
    def fill_frame(self) -> bool:
        if self.__fill_frame is None:
            settings = carb.settings.get_settings()
            width, height = settings.get('/app/renderer/resolution/width') or 0, settings.get('/app/renderer/resolution/height') or 0
            self.__fill_frame = (width <= 0) or (height <= 0)
        return self.__fill_frame

    @fill_frame.setter
    def fill_frame(self, value: bool):
        settings = carb.settings.get_settings()
        legacy_window = self.legacy_window
        self.__fill_frame = bool(value)
        if self.__fill_frame:
            self.__saved_resolution = legacy_window.get_texture_resolution()
            legacy_window.set_texture_resolution(-1, -1)
            settings.set('/app/renderer/resolution/width', -1)
            settings.set('/app/renderer/resolution/height', -1)
        else:
            if self.__saved_resolution is None:
                width, height = settings.get('/app/renderer/resolution/width') or 0, settings.get('/app/renderer/resolution/height') or 0
                self.__saved_resolution = (width, height) if (width > 0 and height > 0) else (1280, 720)
            legacy_window.set_texture_resolution(self.__saved_resolution[0], self.__saved_resolution[1])

    @property
    def fps(self) -> float:
        '''Return the frames-per-second this Viewport is running at'''
        return self.legacy_window.get_fps()

    @property
    def id(self):
        '''Return a hashable value for the Viewport'''
        return self.legacy_window.get_id()

    @property
    def frame_info(self):
        return { 'viewport_handle' : self.legacy_window.get_id() }

    @property
    def usd_context_name(self) -> str:
        '''Return the name of the omni.usd.UsdContext this Viewport is attached to'''
        return self.legacy_window.get_usd_context_name()

    @property
    def usd_context(self):
        '''Return the omni.usd.UsdContext this Viewport is attached to'''
        return omni.usd.get_context(self.usd_context_name)

    @property
    def stage(self):
        '''Return the Usd.Stage of the omni.usd.UsdContext this Viewport is attached to'''
        return self.usd_context.get_stage()

    @property
    def projection(self):
        '''Return the projection of the UsdCamera in terms of the ui element it sits in.'''
        # Check if there are cached values from the ScenView subscriptions
        legacy_subs = _LegacySceneView.get(self.__window_name)
        if legacy_subs:
            projection = legacy_subs.projection
            if projection:
                return projection

        # Get the info from USD
        stage = self.stage
        cam_prim = stage.GetPrimAtPath(self.camera_path) if stage else None
        if cam_prim:
            usd_camera = UsdGeom.Camera(cam_prim)
            if usd_camera:
                image_aspect, canvas_aspect = self._aspect_ratios
                return self._conform_projection(self._conform_policy(), usd_camera, image_aspect, canvas_aspect)

        return Gf.Matrix4d(1)

    @property
    def transform(self) -> Gf.Matrix4d:
        '''Return the world-space transform of the UsdGeom.Camera being used to render'''
        # Check if there are cached values from the ScenView subscriptions
        legacy_subs = _LegacySceneView.get(self.__window_name)
        if legacy_subs:
            view = legacy_subs.view
            if view:
                return view.GetInverse()

        # Get the info from USD
        stage = self.stage
        cam_prim = stage.GetPrimAtPath(self.camera_path) if stage else None
        if cam_prim:
            imageable = UsdGeom.Imageable(cam_prim)
            if imageable:
                return imageable.ComputeLocalToWorldTransform(self.time)

        return Gf.Matrix4d(1)

    @property
    def view(self) -> Gf.Matrix4d:
        '''Return the inverse of the world-space transform of the UsdGeom.Camera being used to render'''
        # Check if there are cached values from the ScenView subscriptions
        legacy_subs = _LegacySceneView.get(self.__window_name)
        if legacy_subs:
            view = legacy_subs.view
            if view:
                return view

        # Get the info from USD
        return self.transform.GetInverse()

    @property
    def world_to_ndc(self):
        return self.view * self.projection

    @property
    def ndc_to_world(self):
        return self.world_to_ndc.GetInverse()

    @property
    def time(self) -> Usd.TimeCode:
        '''Return the Usd.TimeCode this Viewport is using'''
        if self.__timeline is None:
            self.__timeline = omni.timeline.get_timeline_interface()

        time = self.__timeline.get_current_time()
        return Usd.TimeCode(omni.usd.get_frame_time_code(time, self.stage.GetTimeCodesPerSecond()))

    def map_ndc_to_texture(self, mouse: Sequence[float]):
        image_aspect, canvas_aspect = self._aspect_ratios
        if image_aspect < canvas_aspect:
            ratios = (image_aspect / canvas_aspect, 1)
        else:
            ratios = (1, canvas_aspect / image_aspect)

        # Move into viewport's NDC-space: [-1, 1] bound by viewport
        mouse = (mouse[0] / ratios[0], mouse[1] / ratios[1])
        # Move from NDC space to texture-space [-1, 1] to [0, 1]
        def check_bounds(coord): return coord >= -1 and coord <= 1
        return tuple((x + 1.0) * 0.5 for x in mouse), self if (check_bounds(mouse[0]) and check_bounds(mouse[1])) else None

    def map_ndc_to_texture_pixel(self, mouse: Sequence[float]):
        # Move into viewport's uv-space: [0, 1]
        mouse, viewport = self.map_ndc_to_texture(mouse)
        # Then scale by resolution flipping-y
        resolution = self.resolution
        return (int(mouse[0] * resolution[0]), int((1.0 - mouse[1]) * resolution[1])), viewport

    def request_query(self, pixel: Sequence[int], callback: Callable, *args):
        # TODO: This can't be made 100% compatible without changes to legacy Viewport
        # New Viewport will query based on a pixel co-ordinate, regardless of mouse-state; with object and world-space-position
        # Legacy Viewport will invoke callback on mouse-click only; with world-space-position info only
        self.legacy_window.query_next_picked_world_position(lambda pos: callback(True, pos) if pos else callback(False, None))
        carb.log_warn("request_query not full implemented, will only receive a position from a Viewport click")

    @property
    def hydra_engine(self) -> str:
        '''Get the name of the active omni.hydra.engine for this Viewport'''
        return self.legacy_window.get_active_hydra_engine()

    @hydra_engine.setter
    def hydra_engine(self, hd_engine: str) -> str:
        '''Set the name of the active omni.hydra.engine for this Viewport'''
        return self.set_hd_engine(hd_engine)

    @property
    def render_mode(self):
        '''Get the render-mode for the active omni.hydra.engine used in this Viewport'''
        hd_engine = self.hydra_engine
        if bool(hd_engine):
            return self.__settings.get(self.__render_mode_setting(hd_engine))

    @render_mode.setter
    def render_mode(self, render_mode: str):
        '''Set the render-mode for the active omni.hydra.engine used in this Viewport'''
        hd_engine = self.hydra_engine
        if bool(hd_engine):
            self.__settings.set_string(self.__render_mode_setting(hd_engine), str(render_mode))

    def set_hd_engine(self, hd_engine: str, render_mode: str = None):
        '''Set the active omni.hydra.engine for this Viewport, and optionally its render-mode'''
        # self.__settings.set_string("/renderer/active", hd_engine)
        if bool(hd_engine) and bool(render_mode):
            self.__settings.set_string(self.__render_mode_setting(hd_engine), render_mode)
        self.legacy_window.set_active_hydra_engine(hd_engine)

    async def wait_for_rendered_frames(self, additional_frames: int = 0) -> bool:
        '''Asynchrnously wait until the renderer has delivered an image'''
        app = omni.kit.app.get_app_interface()
        legacy_window = self.legacy_window
        if legacy_window:
            viewport_ldr_rp = None
            while viewport_ldr_rp is None:
                await app.next_update_async()
                viewport_ldr_rp = legacy_window.get_drawable_ldr_resource()

        while additional_frames > 0:
            additional_frames = additional_frames - 1
            await app.next_update_async()

        return True

    def add_scene_view(self, scene_view):
        window_name = self.__window_name
        legacy_subs = _LegacySceneView.get(self.__window_name, self.legacy_window)
        if not legacy_subs:
            raise RuntimeError('Could not create _LegacySceneView subscriptions')
        legacy_subs.add_scene_view(scene_view, self.projection, self.view)

    def remove_scene_view(self, scene_view):
        if not scene_view:
            raise RuntimeError('Provided scene_view is invalid')
        _LegacySceneView.remove_scene_view_for_window(self.__window_name, scene_view)

    ### Everything below is NOT part of the new Viewport-API

    @property
    def legacy_window(self):
        '''Expose the underlying legacy viewport for access if needed (not a real ViewportAPI method)'''
        import omni.kit.viewport_legacy as vp_legacy
        vp_iface = vp_legacy.get_viewport_interface()
        return vp_iface.get_viewport_window(vp_iface.get_instance(self.__window_name))

    @property
    def _aspect_ratios(self) -> Tuple[float]:
        return _LegacySceneView.aspect_ratios(self.__window_name, self.legacy_window)

    def _conform_projection(self, policy, camera: UsdGeom.Camera, image_aspect: float, canvas_aspect: float, projection: Sequence[float] = None):
        '''For the given camera (or possible incoming projection) return a projection matrix that matches the rendered image
        but keeps NDC co-ordinates for the texture bound to [-1, 1]'''
        if projection is None:
            # If no projection is provided, conform the camera based on settings
            # This wil adjust apertures on the gf_camera
            gf_camera = camera.GetCamera(self.time)
            if policy == CameraUtil.DontConform:
                # For DontConform, still have to conform for the final canvas
                if image_aspect < canvas_aspect:
                    gf_camera.horizontalAperture = gf_camera.horizontalAperture * (canvas_aspect / image_aspect)
                else:
                    gf_camera.verticalAperture = gf_camera.verticalAperture * (image_aspect / canvas_aspect)
            else:
                CameraUtil.ConformWindow(gf_camera, policy, image_aspect)

            projection = gf_camera.frustum.ComputeProjectionMatrix()
        else:
            projection = Gf.Matrix4d(*projection)

        # projection now has the rendered image projection
        # Conform again based on canvas size so projection extends with the Viewport sits in the UI
        if image_aspect < canvas_aspect:
            policy2 = CameraUtil.MatchVertically
        else:
            policy2 = CameraUtil.MatchHorizontally

        if policy != CameraUtil.DontConform:
            projection = CameraUtil.ConformedWindow(projection, policy2, canvas_aspect)

        return projection

    def _conform_policy(self):
        conform_setting = self.__settings.get("/app/hydra/aperture/conform")
        if (conform_setting is None) or (conform_setting == 1) or (conform_setting == 'horizontal'):
            return CameraUtil.MatchHorizontally
        if (conform_setting == 0) or (conform_setting == 'vertical'):
            return CameraUtil.MatchVertically
        if (conform_setting == 2) or (conform_setting == 'fit'):
            return CameraUtil.Fit
        if (conform_setting == 3) or (conform_setting == 'crop'):
            return CameraUtil.Crop
        if (conform_setting == 4) or (conform_setting == 'stretch'):
            return CameraUtil.DontConform
        return CameraUtil.MatchHorizontally

    def __render_mode_setting(self, hd_engine: str) -> str:
        return f'/{hd_engine}/rendermode' if hd_engine != 'iray' else '/rtx/iray/rendermode'


class _LegacySceneView:
    __g_legacy_subs = {}

    @staticmethod
    def aspect_ratios(window_name: str, legacy_window):
        canvas_aspect = 1
        viewport_window = omni.ui.Workspace.get_window(window_name)
        if viewport_window:
            # XXX: Why do some windows have no frame !?
            if hasattr(viewport_window, 'frame'):
                canvas = viewport_window.frame
                canvas_size = (canvas.computed_width, canvas.computed_height)
            else:
                canvas_size = (1, 1)
            canvas_aspect = canvas_size[0] / canvas_size[1] if canvas_size[1] else 1

        resolution = legacy_window.get_texture_resolution()
        image_aspect = resolution[0] / resolution[1] if resolution[1] else 1
        return image_aspect, canvas_aspect

    @staticmethod
    def get(window_name: str, legacy_window = None):
        legacy_subs = _LegacySceneView.__g_legacy_subs.get(window_name)
        if legacy_subs is None and legacy_window:
            legacy_subs = _LegacySceneView(window_name, legacy_window)
            _LegacySceneView.__g_legacy_subs[window_name] = legacy_subs
        return legacy_subs

    @staticmethod
    def remove_scene_view_for_window(window_name: str, scene_view):
        if not scene_view:
            raise RuntimeError('Provided scene_view is invalid')
        legacy_subs = _LegacySceneView.__g_legacy_subs.get(window_name)
        if not legacy_subs:
            raise RuntimeError('Removing a SceneView for Viewport that is not trakcing any')
        if legacy_subs.remove_scene_view(scene_view):
            return
        legacy_subs.destroy()
        del _LegacySceneView.__g_legacy_subs[window_name]

    def __init__(self, window_name: str, legacy_window):
        self.__window_name = window_name
        self.__draw_sub, self.__update_sub = None, None
        self.__scene_views = []
        self.__projection = None
        self.__view = None

        events = legacy_window.get_ui_draw_event_stream()
        self.__draw_sub = events.create_subscription_to_pop(self.__on_draw, name=f'omni.kit.viewport.utility.LegacyViewportAPI.{window_name}.draw')
        events = omni.kit.app.get_app().get_update_event_stream()
        self.__update_sub = events.create_subscription_to_pop(self.__on_update, name='omni.kit.viewport.utility.LegacyViewportAPI.{window_name}.update')

    @property
    def projection(self) -> Gf.Matrix4d:
        # Return a copy as the reference would be mutable
        return Gf.Matrix4d(self.__projection) if self.__projection else None

    @property
    def view(self) -> Gf.Matrix4d:
        # Return a copy as the reference would be mutable
        return Gf.Matrix4d(self.__view) if self.__view else None

    def __del__(self):
        self.destroy()

    def destroy(self):
        self.__scene_views = []
        if self.__draw_sub:
            self.__draw_sub = None
        if self.__update_sub:
            self.__update_sub = None
        self.__save_position = None
        self.__scene_views = []

    @staticmethod
    def _flatten_matrix(m: Gf.Matrix4d) -> List[float]:
        m0, m1, m2, m3 = m[0], m[1], m[2], m[3]
        return [m0[0], m0[1], m0[2], m0[3], m1[0], m1[1], m1[2], m1[3], m2[0], m2[1], m2[2], m2[3], m3[0], m3[1], m3[2], m3[3]]

    def __get_legacy_window(self):
        import omni.kit.viewport_legacy as vp_legacy
        vp_iface = vp_legacy.get_viewport_interface()
        window_name = self.__window_name
        viewport_handle = vp_iface.get_instance(window_name)
        viewport_window = vp_iface.get_viewport_window(viewport_handle)
        if not viewport_window:
            self.destroy()
            del _LegacySceneView.__g_legacy_subs[window_name]
        return window_name, viewport_window

    def __on_update(self, *args):
        window_name, legacy_window = self.__get_legacy_window()
        visible = legacy_window.is_visible() if legacy_window else False
        viewport_window = omni.ui.Workspace.get_window(window_name)
        if viewport_window and (visible != viewport_window.visible):
            pruned_views = []
            for sv in self.__scene_views:
                scene_view = sv()
                if scene_view:
                    scene_view.visible = visible
                    pruned_views.append(sv)
            self.__scene_views = pruned_views

            omni.ui.Workspace.show_window(window_name, visible)
            # viewport_window.visible = visible
            # if not visible:
            #     self.__save_position = viewport_window.position_x, viewport_window.position_y
            #     viewport_window.position_x = omni.ui.Workspace.get_main_window_width()
            #     viewport_window.position_y = omni.ui.Workspace.get_main_window_height()
            # elif self.__save_position:
            #     viewport_window.position_x = self.__save_position[0]
            #     viewport_window.position_y = self.__save_position[1]
            #     self.__save_position = None

    def __on_draw(self, event, *args):
        vm = event.payload["viewMatrix"]
        pm = event.payload["projMatrix"]

        # Conform the Projection to the layout of the texture in the Window
        window_name, legacy_window = self.__get_legacy_window()
        image_aspect, canvas_aspect = _LegacySceneView.aspect_ratios(window_name, legacy_window)
        if image_aspect < canvas_aspect:
            policy = CameraUtil.MatchVertically
        else:
            policy = CameraUtil.MatchHorizontally

        proj_matrix = Gf.Matrix4d(pm[0], pm[1], pm[2], pm[3], pm[4], pm[5], pm[6], pm[7], pm[8], pm[9], pm[10], pm[11], pm[12], pm[13], pm[14], pm[15])
        # Fix up RTX projection Matrix for omni.ui.scene...ultimately this should be passed without any compensation
        if pm[15] == 0.0:
            proj_matrix = Gf.Matrix4d(1.0, 0.0, -0.0, 0.0, 0.0, 1.0, 0.0, -0.0, -0.0, 0.0, 1.0, -1.0, 0.0, -0.0, 0.0, -2.0) * proj_matrix
        proj_matrix = CameraUtil.ConformedWindow(proj_matrix, policy, canvas_aspect)

        # Cache these for lokkup later
        self.__projection = proj_matrix
        self.__view = Gf.Matrix4d(vm[0], vm[1], vm[2], vm[3], vm[4], vm[5], vm[6], vm[7], vm[8], vm[9], vm[10], vm[11], vm[12], vm[13], vm[14], vm[15])

        # Flatten into list for model
        pm = _LegacySceneView._flatten_matrix(proj_matrix)

        pruned_views = []
        for sv in self.__scene_views:
            scene_view = sv()
            if scene_view:
                model = scene_view.model
                model.set_floats("projection", pm)
                model.set_floats("view", vm)
                pruned_views.append(sv)
        self.__scene_views = pruned_views

    def add_scene_view(self, scene_view, projection: Gf.Matrix4d, view: Gf.Matrix4d):
        # Sync the model once now
        model = scene_view.model
        model.set_floats("projection", _LegacySceneView._flatten_matrix(projection))
        model.set_floats("view", _LegacySceneView._flatten_matrix(view))
        # And save it for subsequent synchrnoizations
        import weakref
        self.__scene_views.append(weakref.ref(scene_view))

    def remove_scene_view(self, scene_view) -> bool:
        for sv in self.__scene_views:
            if sv() == scene_view:
                self.__scene_views.remove(sv)
                break
        # Return whether this sub is needed anymore
        return True if self.__scene_views else False
