import os
import os.path
import omni.kit.test
import omni.usd
import carb
import carb.settings
import carb.tokens
import pathlib
import gc
import unittest


from omni.kit.capture.viewport import CaptureOptions, CaptureExtension
from omni.kit.viewport.utility import get_active_viewport
from omni.kit.test_suite.helpers import wait_stage_loading
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")
TEST_DIR = str(pathlib.Path(__file__).parent.parent.parent.parent.parent.parent)


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

    # Before running each test
    async def setUp(self):
        self._usd_context = ''
        self._frames_per_sec_to_wait_for_capture_resource_ready = 60
        test_usd = TEST_DIR + "/data/tests/usd/PTAccumulateTest.usd"
        carb.log_warn(f"testing capture with usd {test_usd}")
        self._capture_folder_1spp = "rp_1spp"
        self._capture_folder_32spp = "rp_32spp"
        self._capture_folder_sequence = "rp_seq"
        self._context = omni.usd.get_context()
        self._capture_instance = CaptureExtension().get_instance()
        await self._context.open_stage_async(test_usd)
        await wait_stage_loading()

        for _ in range(2):
            await omni.kit.app.get_app().next_update_async()

    async def tearDown(self) -> None:
        self._capture_instance = None

    async def _test_render_product_capture(self, image_folder, render_product_name, spp, enable_hdr=True, sequence_capture=False, start_frame=0, end_frame=10):
        viewport_api = get_active_viewport(self._usd_context)
        # Wait until the viewport has valid resources
        await viewport_api.wait_for_rendered_frames()

        # wait for PT to resolve after loading, otherwise it's easy to crash
        i = self._frames_per_sec_to_wait_for_capture_resource_ready * 50
        while i > 0:
            await omni.kit.app.get_app().next_update_async()
            i -= 1

        # assert True
        capture_folder_name = image_folder
        filePath = pathlib.Path(OUTPUTS_DIR).joinpath(capture_folder_name)
        options = CaptureOptions()
        options.file_type = ".exr"
        options.output_folder = str(filePath)
        make_sure_directory_existed(options.output_folder)
        clean_files_in_directory(options.output_folder, ".exr")
        options.hdr_output = enable_hdr
        options.camera = viewport_api.camera_path.pathString
        options.render_preset = omni.kit.capture.viewport.CaptureRenderPreset.PATH_TRACE
        options.render_product = render_product_name
        options.path_trace_spp = spp
        # to reduce memory consumption and save time
        options.res_width = 600
        options.res_height = 337

        if sequence_capture:
            options.capture_every_Nth_frames = 1
            options.start_frame = start_frame
            options.end_frame = end_frame

        self._capture_instance.options = options
        self._capture_instance.start()
        options = None

    async def _verify_file_exist(self, capture_folder_name, aov_channel):
        exr_path = self._get_exr_path(capture_folder_name, aov_channel)
        carb.log_warn(f"Verifying image: {exr_path}")
        await wait_for_image_writing(exr_path, 40)
        assert os.path.isfile(exr_path)

    def _get_exr_path(self, capture_folder_name, aov_channel):
        image_folder = pathlib.Path(OUTPUTS_DIR).joinpath(capture_folder_name)
        exr_path = os.path.join(str(image_folder), "Capture1_" + aov_channel + ".exr")
        return exr_path

    def _get_captured_image_size(self, capture_folder_name, aov_channel):
        image_folder = pathlib.Path(OUTPUTS_DIR).joinpath(capture_folder_name)
        exr_path = os.path.join(str(image_folder), "Capture1_" + aov_channel + ".exr")
        if os.path.isfile(exr_path):
            return os.path.getsize(exr_path)
        else:
            return 0

    async def test_capture_1_default_rp(self):
        if True:
            # skip the test for vp2 due to OM-77175: we know it will fail and have a bug for it so mark it true here
            # and re-enable it after the bug gets fixed

            image_folder_name = "vp2_default_rp"
            default_rp_name = "/Render/RenderProduct_omni_kit_widget_viewport_ViewportTexture_0"

            assert True
            for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 10):
                await omni.kit.app.get_app().next_update_async()
            return

        await self._test_render_product_capture(image_folder_name, default_rp_name, 16, enable_hdr=False)
        exr_path = self._get_exr_path(image_folder_name, "LdrColor")
        await wait_for_image_writing(exr_path)
        await self._test_render_product_capture(image_folder_name, default_rp_name, 16)
        await self._verify_file_exist(image_folder_name, "LdrColor")
        for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 10):
            await omni.kit.app.get_app().next_update_async()

    async def test_capture_2_user_created_rp(self):
        await self._test_render_product_capture(self._capture_folder_1spp, "/Render/RenderView", 1)
        for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 10):
            await omni.kit.app.get_app().next_update_async()
        await self._verify_file_exist(self._capture_folder_1spp, "LdrColor")
        await self._verify_file_exist(self._capture_folder_1spp, "HdrColor")
        await self._verify_file_exist(self._capture_folder_1spp, "PtGlobalIllumination")
        await self._verify_file_exist(self._capture_folder_1spp, "PtReflections")
        await self._verify_file_exist(self._capture_folder_1spp, "PtWorldNormal")

        # the render product capture process will write the image file multiple times during the capture
        # so give it more cycles to settle down to check file size
        for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 20):
            await omni.kit.app.get_app().next_update_async()

    async def test_capture_3_rp_accumulation(self):
        await self._test_render_product_capture(self._capture_folder_32spp, "/Render/RenderView", 32)
        for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 10):
            await omni.kit.app.get_app().next_update_async()
        await self._verify_file_exist(self._capture_folder_32spp, "LdrColor")

        # the render product capture process will write the image file multiple times during the capture
        # so give it more cycles to settle down to check file size
        for _ in range(self._frames_per_sec_to_wait_for_capture_resource_ready * 20):
            await omni.kit.app.get_app().next_update_async()

        # compare result of test_capture_2_user_created_rp images to check accumulation works or not
        # if it works, more spp will result in smaller file size
        file_size_1spp = self._get_captured_image_size(self._capture_folder_1spp, "LdrColor")
        file_size_32spp = self._get_captured_image_size(self._capture_folder_32spp, "LdrColor")
        carb.log_warn(f"File size of 1spp capture is {file_size_1spp} bytes, and file size of 64spp capture is {file_size_32spp}")
        expected_result = file_size_32spp != file_size_1spp
        assert expected_result

    async def test_capture_4_rp_sequence(self):
        image_folder = pathlib.Path(OUTPUTS_DIR).joinpath(self._capture_folder_sequence, "Capture_frames")
        clean_files_in_directory(image_folder, ".exr")
        await self._test_render_product_capture(self._capture_folder_sequence, "/Render/RenderView", 1, True, True, 0, 10)
        start_frame_exr_path_1 = os.path.join(str(image_folder), "Capture.0000_" + "LdrColor" + ".exr")
        start_frame_exr_path_2 = os.path.join(str(image_folder), "Capture.0000_" + "PtGlobalIllumination" + ".exr")
        end_frame_exr_path_1 = os.path.join(str(image_folder), "Capture.0010_" + "LdrColor" + ".exr")
        end_frame_exr_path_2 = os.path.join(str(image_folder), "Capture.0010_" + "PtGlobalIllumination" + ".exr")
        await wait_for_image_writing(end_frame_exr_path_2, 50)
        assert os.path.isfile(start_frame_exr_path_1)
        assert os.path.isfile(start_frame_exr_path_2)
        assert os.path.isfile(end_frame_exr_path_1)
        assert os.path.isfile(end_frame_exr_path_2)
