# Copyright (c) 2018-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__ = ["DocumentationBuilderCapture"]

import asyncio

import omni.appwindow
import omni.kit.app
import omni.kit.imgui_renderer
import omni.kit.renderer.bind
import omni.renderer_capture
import omni.ui as ui


class DocumentationBuilderCapture:
    r"""
    Captures the omni.ui widgets to the image file. Works like this:

        with Capture("c:\temp\1.png", 600, 600):
            ui.Button("Hello World")
    """

    def __init__(self, path: str, width: int, height: int, window_name: str = "Capture Window"):
        self._path: str = path
        self._dpi = ui.Workspace.get_dpi_scale()
        self._width: int = int(width * self._dpi)
        self._height: int = int(height * self._dpi)
        self._window_name: str = window_name

        # Interfaces
        self._app_window_factory = omni.appwindow.acquire_app_window_factory_interface()
        self._renderer = omni.kit.renderer.bind.acquire_renderer_interface()
        self._imgui_renderer = omni.kit.imgui_renderer.acquire_imgui_renderer_interface()
        self._renderer_capture = omni.renderer_capture.acquire_renderer_capture_interface()

        self._virtual_app_window = None
        self._window = None

    def __enter__(self):
        # Create virtual OS window
        self._virtual_app_window = self._app_window_factory.create_window_by_type(omni.appwindow.WindowType.VIRTUAL)
        self._virtual_app_window.startup_with_desc(
            title=self._window_name,
            width=self._width,
            height=self._height,
            scale_to_monitor=False,
            dpi_scale_override=1.0,
        )
        self._imgui_renderer.attach_app_window(self._virtual_app_window)

        # Create omni.ui window
        self._window = ui.Window(self._window_name, flags=ui.WINDOW_FLAGS_NO_TITLE_BAR, padding_x=0, padding_y=0)
        self._window.move_to_app_window(self._virtual_app_window)
        self._window.width = self._width
        self._window.height = self._height

        # Enter the window frame
        self._window.frame.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Exit the window frame
        self._window.frame.__exit__(exc_type, exc_val, exc_tb)

        # Capture and detach at the next frame
        asyncio.ensure_future(
            DocumentationBuilderCapture.__capture_detach(
                self._renderer_capture, self._path, self._renderer, self._virtual_app_window, self._window
            )
        )

    @staticmethod
    async def __capture_detach(renderer_capture, path, renderer, iappwindow, window):
        # The next frame
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

        # Capture
        renderer_capture.capture_next_frame_texture(path, renderer.get_framebuffer_texture(iappwindow), iappwindow)

        await omni.kit.app.get_app().next_update_async()

        # Detach
        renderer.detach_app_window(iappwindow)
