import os
import os.path
import omni.kit.test
import carb
import carb.settings
import carb.tokens
import asyncio

import omni.ui as ui

from omni.kit.capture.viewport import CaptureOptions, CaptureExtension, CaptureRenderPreset
from omni.kit.viewport.utility import get_active_viewport, create_viewport_window, get_active_viewport_and_window

async def wait_n_updates(n_frames: int = 3):
    app = omni.kit.app.get_app()
    for _ in range(n_frames):
        await app.next_update_async()

async def wait_for_image_writing(image_path, seconds_to_wait: float = 30, seconds_to_sleep: float = 0.5):
    seconds_tried = 0.0
    carb.log_info(f"Waiting for {image_path} to be written to disk.")
    while seconds_tried < seconds_to_wait:
        if os.path.isfile(image_path) and os.access(image_path, os.R_OK):
            carb.log_info(f"{image_path} is written to disk in {seconds_tried}s.")
            break
        else:
            await asyncio.sleep(seconds_to_sleep)
            seconds_tried += seconds_to_sleep

    if seconds_tried >= seconds_to_wait:
        carb.log_warn(f"Waiting for {image_path} timed out..")


def clean_files_in_directory(directory, suffix):
    if not os.path.exists(directory):
        return
    images = os.listdir(directory)
    for item in images:
        if item.endswith(suffix):
            os.remove(os.path.join(directory, item))


def make_sure_directory_existed(directory):
    if not os.path.exists(directory):
        try:
            os.makedirs(directory, exist_ok=True)
        except OSError as error:
            carb.log_warn(f"Directory cannot be created: {dir}")
            return False
    return True



async def capture_png(camera_path, frame, output_folder, output_filename):
    settings = carb.settings.get_settings()
    _capture_instance = CaptureExtension().get_instance()
    await wait_n_updates()


    viewport_api = get_active_viewport()

    if not viewport_api:
        raise RuntimeError("No active Viewport")

    # Set the Viewport's active camera to the
    # camera prim path you want to switch to.
    viewport_api.camera_path = camera_path
  

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

    
    options = CaptureOptions()
    options.file_type = ".png"
    
    # OM-112018 happens in PT captures
    options.render_preset = CaptureRenderPreset.PATH_TRACE
    
    options.real_time_settle_latency_frames = 0
    options.path_trace_spp = 16
    options.output_folder = output_folder
    options.file_name = output_filename
    options.overwrite_existing_frames = True
    make_sure_directory_existed(options.output_folder)
    
    clean_files_in_directory(output_folder, ".png")

    png_path = os.path.join(output_folder, output_filename + "1.png")
    
    options.hdr_output = False
    options.camera = viewport_api.camera_path.pathString
    carb.log_warn("vp_cam_path: "+viewport_api.camera_path.pathString)

    _capture_instance.options = options
    _capture_instance.start()

    await wait_for_image_writing(png_path, seconds_to_wait=30)    
    carb.log_warn(os.path.isfile(png_path))

    _capture_instance.cancel()
    await wait_n_updates(20)



output_folder = r"D:\temp\ov_testrendering"

camera_paths = [
    "/cam_ABCD",
    "/cam_DCBA"    
]

for camera_path in camera_paths:
    carb.log_warn(camera_path)
    asyncio.ensure_future(capture_png(camera_path, 0, output_folder, camera_path.split("/")[-1]))