from typing import Type
import os.path
import omni.kit.test
import carb
import carb.settings
import carb.tokens
import pathlib

import omni.kit.capture.capture_options as _capture_options
from omni.kit.viewport.utility import get_active_viewport, capture_viewport_to_file, create_viewport_window

OUTPUTS_DIR = omni.kit.test.get_test_output_path()

class TestCaptureHdr(omni.kit.test.AsyncTestCase):

    # Before running each test
    async def setUp(self):
        settings = carb.settings.get_settings()
        settings.set("/app/asyncRendering", False)
        settings.set("/app/asyncRenderingLowLatency", False)
        settings.set("/app/captureFrame/hdr", True)

        # Setup viewport
        self._usd_context = ''
        await omni.usd.get_context(self._usd_context).new_stage_async()

    async def test_hdr_capture(self):
        viewport_api = get_active_viewport(self._usd_context)
        if viewport_api is None:
            return

        # Wait until the viewport has valid resources
        await viewport_api.wait_for_rendered_frames()

        capture_filename = "capture.hdr_test.exr"
        filePath = pathlib.Path(OUTPUTS_DIR).joinpath(capture_filename)

        capture_viewport_to_file(viewport_api, str(filePath))

        await omni.kit.app.get_app_interface().next_update_async()
        await omni.kit.app.get_app_interface().next_update_async()
        await omni.kit.app.get_app_interface().next_update_async()
        assert os.path.isfile(str(filePath))

    # Make sure we do not crash in the unsupported multi-view case
    async def do_test_hdr_multiview_capture(self, test_legacy: bool):
        viewport_api = get_active_viewport(self._usd_context)
        if viewport_api is None:
            return

        # Wait until the viewport has valid resources
        await viewport_api.wait_for_rendered_frames()

        second_vp_window = create_viewport_window('Viewport 2', width=256, height=256)
        await second_vp_window.viewport_api.wait_for_rendered_frames()

        capture_filename = "capture.hdr_test.multiview.exr"
        filePath = pathlib.Path(OUTPUTS_DIR).joinpath(capture_filename)

        # Multiview HDR output is not yet fully supported, confirm this exits gracefully
        capture_viewport_to_file(viewport_api, str(filePath))

        await omni.kit.app.get_app_interface().next_update_async()
        await omni.kit.app.get_app_interface().next_update_async()
        await omni.kit.app.get_app_interface().next_update_async()
        assert os.path.isfile(str(filePath))

        second_vp_window.destroy()
        del second_vp_window

    async def test_hdr_multiview_capture_legacy(self):
        await self.do_test_hdr_multiview_capture(True)

    async def test_hdr_multiview_capture(self):
        await self.do_test_hdr_multiview_capture(False)
