# Copyright (c) 2023, 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__ = ["enable_viewport_rendering", "start_render_loading_ui"]

from typing import Optional
import omni.usd


_g_enabled_renderers = []
g_renderers_ready = []

def start_render_loading_ui(ext_manager: "omni.ext.ExtensionManager", renderer: str):
    import carb
    import time

    time_begin = time.time()

    # Extension to show viewport loading progress
    ext_manager.set_extension_enabled_immediate("omni.activity.profiler", True)
    ext_manager.set_extension_enabled_immediate("omni.activity.pump", True)

    carb.settings.get_settings().set("/exts/omni.kit.viewport.ready/startup/enabled", False)
    ext_manager.set_extension_enabled_immediate("omni.kit.viewport.ready", True)

    from omni.kit.viewport.ready.viewport_ready import ViewportReady, ViewportReadyDelegate
    class TimerDelegate(ViewportReadyDelegate):
        def __init__(self, time_begin):
            super().__init__()
            self.__timer_start = time.time()
            self.__vp_ready_cost = self.__timer_start - time_begin

        @property
        def font_size(self) -> float:
            return 48

        @property
        def message(self) -> str:
            rtx_mode = {
                "RaytracedLighting": "Real-Time",
                "PathTracing": "Interactive (Path Tracing)",
                "LightspeedAperture": "Aperture (Game Path Tracer)"
            }.get(carb.settings.get_settings().get("/rtx/rendermode"), "Real-Time")

            renderer_label = {
                "rtx": f"RTX - {rtx_mode}",
                "iray": "RTX - Accurate (Iray)",
                "pxr": "Pixar Storm",
                "index": "RTX - Scientific (IndeX)"
            }.get(renderer, renderer)

            return f"Waiting for {renderer_label} to start"

        def on_complete(self):
            global g_renderers_ready
            g_renderers_ready.append(renderer)
            rtx_load_time = time.time() - self.__timer_start
            super().on_complete()
            carb.log_info(f"Time until pixel: {rtx_load_time}")
            carb.log_info(f"ViewportReady cost: {self.__vp_ready_cost}")

            settings = carb.settings.get_settings()
            if settings.get("/exts/omni.kit.welcome.window/autoSetupOnClose"):
                async def async_enable_all_renderers(app, renderer: str, all_renderers: str):
                    for enabled_renderer in all_renderers.split(","):
                        if enabled_renderer != renderer:
                            enable_viewport_rendering(enabled_renderer, False)
                            await app.next_update_async()

                import omni.kit.app
                import omni.kit.async_engine
                omni.kit.async_engine.run_coroutine(async_enable_all_renderers(omni.kit.app.get_app(), renderer,
                                                                               settings.get("/renderer/enabled") or ""))

    return ViewportReady(TimerDelegate(time_begin))


def enable_viewport_rendering(renderer: Optional[str] = None, show_viewport_ready: bool = True):
    import carb
    settings = carb.settings.get_settings()

    # If no renderer is specified, choose the 'active' one or what is set as active, or finally RTX - Realtime
    if renderer is None:
        renderer = (settings.get("/renderer/active") or
                    settings.get("/exts/omni.app.setup/backgroundRendererLoad/renderer") or
                    "rtx")

    # Only enable the Viewport for a renderer once, and only do Viewport Ready the first time
    global _g_enabled_renderers
    if renderer in _g_enabled_renderers:
        if renderer not in g_renderers_ready and show_viewport_ready:
            import omni.kit.app
            app = omni.kit.app.get_app()
            ext_manager = app.get_extension_manager()
            start_render_loading_ui(ext_manager, renderer)
        return
    elif show_viewport_ready and bool(_g_enabled_renderers):
        show_viewport_ready = False
    _g_enabled_renderers.append(renderer)

    settings.set("/renderer/asyncInit", True)
    settings.set("/rtx/hydra/mdlMaterialWarmup", True)

    import omni.kit.app
    app = omni.kit.app.get_app()
    ext_manager = app.get_extension_manager()

    ext_manager.set_extension_enabled_immediate("omni.kit.viewport.bundle", True)
    if renderer == "iray" or renderer == "index":
        ext_manager.set_extension_enabled_immediate(f"omni.hydra.rtx", True)
        ext_manager.set_extension_enabled_immediate(f"omni.hydra.{renderer}", True)
    else:
        ext_manager.set_extension_enabled_immediate(f"omni.kit.viewport.{renderer}", True)

    # we might need to have the engine_name as part of the flow
    ext_manager.set_extension_enabled("omni.rtx.settings.core", True)

    # TODO: It will block omni.kit.viewport.Ready.ViewportReady while it is enabled before this function
    # omni.usd.add_hydra_engine(renderer, usd_context)

    if show_viewport_ready:
        start_render_loading_ui(ext_manager, renderer)
