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

from omni.kit.capture.viewport import CaptureOptions, CaptureExtension
from omni.kit.viewport.utility import get_active_viewport, create_viewport_window, capture_viewport_to_file
from .test_helper import clean_files_in_directory, make_sure_directory_existed, wait_for_image_writing

KIT_ROOT = pathlib.Path(carb.tokens.acquire_tokens_interface().resolve("${kit}")).parent.parent.parent
OUTPUTS_DIR = KIT_ROOT.joinpath("_testoutput/omni.kit.capture.viewport.images")

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

    # Before running each test
    async def setUp(self):
        self._usd_context = ''
        settings = carb.settings.get_settings()
        self._frames_wait_for_capture_resource_ready = 6
        await omni.usd.get_context(self._usd_context).new_stage_async()

    async def test_png_capture(self):
        viewport_api = get_active_viewport(self._usd_context)
        # Wait until the viewport has valid resources
        await viewport_api.wait_for_rendered_frames()

        capture_filename = "capture_png_test"
        filePath = pathlib.Path(OUTPUTS_DIR).joinpath(capture_filename)
        options = CaptureOptions()
        options.file_type = ".png"
        options.output_folder = str(filePath)
        make_sure_directory_existed(options.output_folder)
        clean_files_in_directory(options.output_folder, ".png")
        exr_path = os.path.join(options._output_folder, "Capture1.png")
        carb.log_warn(f"Capture image path: {exr_path}")
        options.hdr_output = False
        options.camera = viewport_api.camera_path.pathString

        capture_instance = CaptureExtension().get_instance()
        capture_instance.options = options
        capture_instance.start()

        await wait_for_image_writing(exr_path)
        options = None
        capture_instance = None
        gc.collect()
        assert os.path.isfile(exr_path)
