# 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__ = ['ViewportReadyProgressDelegate', 'ViewportReadyDelegate', 'ViewportReady']

import omni.ui
import omni.usd
import carb


class ViewportReadyProgressDelegate:
    def __init__(self, settings: carb.settings.ISettings):
        self.__progress = None
        self.__label = None
        self.__shader_count = None
        self.__enabled_mask_id = None
        self.setup_updates(settings)

    def __del__(self):
        self.destroy()

    def setup_updates(self, settings: carb.settings.ISettings):
        """Setup the ViewportReadyProgressDelegate to get progress updates"""
        try:
            import re
            import omni.activity.core
            import omni.activity.profiler

            self.__re_matcher = re.compile("^.*\.hlsl \(UID: [0-9]+\)")
            persisted_shader_count = settings.get("/persistent/exts/omni.kit.viewport.ready/shader_count")

            # Guard against divide by 0 error in activity_event which shouldn't happen but did once somehow.
            self.__shader_count = (0, persisted_shader_count if persisted_shader_count else 1)

            profiler = omni.activity.profiler.get_activity_profiler()
            if profiler:
                self.__enabled_mask_id = profiler.enable_capture_mask(omni.activity.profiler.CAPTURE_MASK_STARTUP)

            activity = omni.activity.core.get_instance()
            self.__activity = (activity, activity.create_callback_to_pop(self.activity_event))

        except (ImportError, RuntimeError):
            required_exts = "omni.activity.profiler and omni.activity.pump"
            carb.log_error(f"{required_exts} must be enabled with /exts/omni.kit.viewport.ready/activity_progress=true")

    def destroy(self):
        self.__activity = None
        progress, self.__progress = self.__progress, None
        if progress:
            progress.destroy()
        label, self.__label = self.__label, None
        if label:
            label.destroy()
        # Save the total number of shaders processed as the progress max hint for the next launch
        if self.__shader_count:
            carb.settings.get_settings().set("/persistent/exts/omni.kit.viewport.ready/shader_count", self.__shader_count[0])
        # Restore the profile mask to whatever it was before
        enabled_mask_id, self.__enabled_mask_id = self.__enabled_mask_id, None
        if enabled_mask_id is not None:
            profiler = omni.activity.profiler.get_activity_profiler()
            profiler.disable_capture_mask(enabled_mask_id)

    def activity_event(self, node, root_node: bool = True):
        """Event handler for omni.activity messages"""
        # This is a rather delicate method relying on strings being constant to filter properly
        node_name = node.name or ""
        if not root_node:
            if self.__re_matcher.search(node_name) is not None:
                self.__shader_count = self.__shader_count[0] + 1, self.__shader_count[1]
                self.set_progress(self.__shader_count[0] / self.__shader_count[1], f"Compiling: {node_name}")
        elif node_name == "Ray Tracing Pipeline":
            for node_id in range(node.child_count):
                self.activity_event(node.get_child(node_id), False)

    def set_progress(self, amount: float, message: str):
        if self.__progress:
            self.__progress.model.set_value(amount)
        if self.__label:
            self.__label.text = message

    def build_ui(self, settings: carb.settings.ISettings, font_size: float, color: omni.ui.color):
        """Build the progress ui with font and color style hints"""
        margin = settings.get("/exts/omni.kit.viewport.ready/activity_progress/margin")
        margin = omni.ui.Percent(20 if margin is None else margin)
        omni.ui.Spacer(width=margin)
        with omni.ui.VStack():
            omni.ui.Spacer(height=5)
            self.__progress = omni.ui.ProgressBar(height=5,
                                                  alignment=omni.ui.Alignment.CENTER_BOTTOM,
                                                  style={"border_width": 2, "border_radius": 5,
                                                         "color": color})
            omni.ui.Spacer(height=5)
            self.__label = omni.ui.Label("", alignment=omni.ui.Alignment.CENTER_TOP, style={"font_size": font_size})
        omni.ui.Spacer(width=margin)


class ViewportReadyDelegate:
    def __init__(self, viewport_name: str = 'Viewport', usd_context_name: str = '', viewport_handle: int = None):
        self.__viewport_name = viewport_name
        self.__usd_context_name = usd_context_name
        self.__viewport_handle = viewport_handle
        self.__frame = None
        # The signal RTX sends that first frame is ready could actually come before omni.ui has called the ui build fn
        self.__destroyed = False

        # Fill in some default based on carb.settings
        settings = carb.settings.get_settings()
        self.__message = settings.get("/exts/omni.kit.viewport.ready/message") or "RTX Loading"
        self.__font_size = settings.get("/exts/omni.kit.viewport.ready/font_size") or 72

        if settings.get("/exts/omni.kit.viewport.ready/activity_progress/enabled"):
            self.__progress = ViewportReadyProgressDelegate(settings)
        else:
            self.__progress = None

    @property
    def message(self) -> str:
        '''Return the string for the default omni.ui.Label.'''
        if self.__message == "RTX Loading":
            active = {
                "rtx": "RTX",
                "iray": "Iray",
                "pxr": "Storm",
                "index": "Index"
            }.get(carb.settings.get_settings().get("/renderer/active"), None)
            if active:
                return f"{active} Loading"

        return self.__message

    @property
    def font_size(self) -> float:
        '''Return the font-size for the default omni.ui.Label.'''
        return self.__font_size

    @property
    def usd_context_name(self) -> str:
        '''Return the omni.usd.UsdContext name this delegate is waiting for completion on.'''
        return self.__usd_context_name

    @property
    def viewport_name(self) -> str:
        '''Return the name of the omni.ui.Window this delegate is waiting for completion on.'''
        return self.__viewport_name

    @property
    def viewport_handle(self) -> int:
        '''Return the ViewportHandle this delegate is waiting for completion on or None for any Viewport.'''
        return self.__viewport_handle

    def on_complete(self):
        '''Callback function invoked when first rendered frame is delivered'''
        # Make sure to always call destroy if any exception is thrown during log_complete
        try:
            self.log_complete()
        finally:
            self.destroy()

    def log_complete(self, msg: str = None, prefix_time_since_start: bool = True):
        '''Callback function that will log to info (and stdout if /app/enableStdoutOutput is enabled)'''
        if msg is None:
            msg = 'RTX ready'

        if prefix_time_since_start:
            import omni.kit.app
            seconds = omni.kit.app.get_app().get_time_since_start_s()
            msg = "[{0:.3f}s] ".format(seconds) + msg

        if carb.settings.get_settings().get('/app/enableStdoutOutput'):
            # Going to sys.stdout directly is cleaner when launched with -info
            import sys
            sys.stdout.write(msg + '\n')
        else:
            carb.log_info(msg)

    def build_label(self) -> omni.ui.Widget:
        '''Simple method to override the basic label showing progress.'''
        return omni.ui.Label(self.message, alignment=omni.ui.Alignment.CENTER, style={"font_size": self.font_size})

    def build_progress(self) -> None:
        settings = carb.settings.get_settings()
        color = settings.get("/exts/omni.kit.viewport.ready/activity_progress/color")
        color = omni.ui.color("#76b900" if color is None else color)
        font_size = settings.get("/exts/omni.kit.viewport.ready/activity_progress/font_size")
        if font_size is None:
            font_size = self.__font_size / 6
        omni.ui.Spacer(height=5)
        with omni.ui.HStack():
            self.__progress.build_ui(settings, font_size, color)

    def build_frame(self) -> None:
        '''Method to override the basic frame showing progress.'''

        # If RTX is ready before omni.ui is calling the build function, then do nothing
        if self.__destroyed:
            return

        self.__frame = omni.ui.ZStack()
        with self.__frame:
            bg_color = carb.settings.get_settings().get('/exts/omni.kit.viewport.ready/background_color')
            if bg_color:
                omni.ui.Rectangle(style={"background_color": bg_color})
            with omni.ui.VStack():
                kwargs = {}
                if self.__progress:
                    kwargs = {"height": omni.ui.Percent(35)}
                omni.ui.Spacer(**kwargs)
                self.build_label()

                if self.__progress:
                    self.build_progress()

                omni.ui.Spacer()

    def build_ui(self) -> None:
        '''Method to build the progress/info state in the Viewport window.'''

        # If RTX is ready before build_ui is called then do nothing
        if self.__destroyed:
            return

        if self.__frame is None:
            viewport_window = omni.ui.Workspace.get_window(self.viewport_name)
            if viewport_window:
                frame = None
                if hasattr(viewport_window, 'get_frame'):
                    frame = viewport_window.get_frame('omni.kit.viewport.ready.ViewportReadyDelegate')
                else:
                    # TODO: Use omni.kit.viewport.utility get_viewport_window to manage omni.ui Wrapper on legacy viewport
                    if hasattr(viewport_window, 'frame'):
                        frame = viewport_window.frame
                if frame:
                    frame.set_build_fn(self.build_frame)

    def __del__(self):
        self.destroy()

    def destroy(self):
        self.__destroyed = True
        progress, self.__progress = self.__progress, None
        if progress:
            progress.destroy()
        if self.__frame:
            self.__frame.visible = False
            if hasattr(self.__frame, 'clear'):
                self.__frame.clear()
            self.__frame.destroy()
            self.__frame = None


class ViewportReady:
    def __init__(self, viewport_ready: ViewportReadyDelegate = None):
        self.__viewport_ready = viewport_ready if viewport_ready else ViewportReadyDelegate()
        usd_context_name = self.__viewport_ready.usd_context_name
        usd_context = omni.usd.get_context(usd_context_name)
        if usd_context:
            # Build the UI objects now
            self.__viewport_ready.build_ui()
            # Subscribe to the eent to wait for frame delivery
            self.__new_frame_sub = usd_context.get_rendering_event_stream().create_subscription_to_push_by_type(
                int(omni.usd.StageRenderingEventType.NEW_FRAME),
                self.__on_event,
                name=f"omni.kit.viewport.ready.ViewportReady"
            )
        else:
            raise RuntimeError(f'No omni.usd.UsdContext found with name "{usd_context_name}"')

    def __del__(self):
        self.destroy()

    def destroy(self):
        self.__new_frame_sub = None
        self.__viewport_ready = None

    def __on_event(self, e: carb.events.IEvent):
        waiting_for = self.__viewport_ready.viewport_handle
        if (waiting_for is None) or (waiting_for == e.payload["viewport_handle"]):
            if self.__viewport_complete():
                # Avoid wrapping callback in try/catch by storing it first, destroying, then calling it
                viewport_ready, self.__viewport_ready = self.__viewport_ready, None
                self.destroy()
                viewport_ready.on_complete()
                del viewport_ready

    def __viewport_complete(self):
        usd_context_name = self.__viewport_ready.usd_context_name
        window_name = self.__viewport_ready.viewport_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(usd_context_name):
                if window.title == window_name:
                    return window.viewport_api.frame_info.get('viewport_handle', None) is not None
        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:
                    return bool(vp_window.get_drawable_ldr_resource() or vp_window.get_drawable_hdr_resource())
        except ImportError:
            pass

        return False
