## Copyright (c) 2021, 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__ = ["TestCaptureAPI"]

import omni.kit.test
from omni.ui.tests.test_base import OmniUiTest
from omni.ui.tests.compare_utils import OUTPUTS_DIR, compare, CompareError

import omni.usd
import omni.ui as ui
from omni.kit.widget.viewport import ViewportWidget
from omni.kit.widget.viewport.capture import *
from omni.kit.test_helpers_gfx.compare_utils import compare as gfx_compare
from omni.kit.test.teamcity import teamcity_publish_image_artifact

import carb
import traceback
from pathlib import Path
import os
from pxr import Usd, UsdRender


CURRENT_PATH = Path(carb.tokens.get_tokens_interface().resolve("${omni.kit.widget.viewport}/data")).absolute().resolve()
TEST_FILES_DIR = CURRENT_PATH.joinpath('tests')
USD_FILES_DIR = TEST_FILES_DIR.joinpath('usd')

TEST_WIDTH, TEST_HEIGHT = 360, 240
NUM_DEFAULT_WINDOWS = 0

# XXX: Make this more accessible in compare_utils
#
def compare_images(image_name, threshold = 10):
    image1 = OUTPUTS_DIR.joinpath(image_name)
    image2 = TEST_FILES_DIR.joinpath(image_name)
    image_diffmap = OUTPUTS_DIR.joinpath(f"{Path(image_name).stem}.diffmap.png")

    gfx_compare(image1, image2, image_diffmap)


class TestByteCapture(ByteCapture):
    def __init__(self, test, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__test = test
        self.__complete = False

    def on_capture_completed(self, buffer, buffer_size, width, height, format):
        self.__test.assertEqual(width, TEST_WIDTH)
        self.__test.assertEqual(height, TEST_HEIGHT)
        self.__test.assertEqual(format, ui.TextureFormat.RGBA8_UNORM)


class TestCaptureAPI(OmniUiTest):
    # Before running each test
    async def setUp(self):
        await super().setUp()

    # After running each test
    async def tearDown(self):
        await super().tearDown()

    async def merged_test_capture_file(self):
        await self.create_test_area(width=TEST_WIDTH, height=TEST_HEIGHT)

        usd_path = str(TEST_FILES_DIR.joinpath('usd/sphere.usda'))
        image_name = f'sphere.png'
        image_path = str(OUTPUTS_DIR.joinpath(image_name))

        usd_context = omni.usd.get_context()
        await usd_context.open_stage_async(usd_path)
        self.assertIsNotNone(usd_context.get_stage())

        vp_window = ui.Window('test_capture_file', width=TEST_WIDTH, height=TEST_HEIGHT, flags=ui.WINDOW_FLAGS_NO_SCROLLBAR)
        self.assertIsNotNone(vp_window)
        with vp_window.frame:
            vp_widget = ViewportWidget(resolution=(TEST_WIDTH, TEST_HEIGHT))
            self.assertIsNotNone(vp_widget)

            capture = vp_widget.viewport_api.schedule_capture(FileCapture(image_path))
            captured_aovs = await capture.wait_for_result()
            self.assertEqual(captured_aovs, ['LdrColor'])
            # Test the frame's metadata is available
            self.assertIsNotNone(capture.view)
            self.assertIsNotNone(capture.projection)
            self.assertEqual(capture.resolution, (TEST_WIDTH, TEST_HEIGHT))
            # Test the viewport_handle is also available for usets that
            self.assertIsNotNone(capture.viewport_handle)

            compare_images(image_name)

        vp_widget.destroy()
        vp_window.destroy()
        del vp_window

    async def merged_test_capture_file_with_format(self):
        await self.create_test_area(width=TEST_WIDTH, height=TEST_HEIGHT)

        usd_path = str(TEST_FILES_DIR.joinpath('usd/sphere.usda'))
        image_name = f'sphere_rle.exr'
        image_path = str(OUTPUTS_DIR.joinpath(image_name))

        usd_context = omni.usd.get_context()
        await usd_context.open_stage_async(usd_path)
        self.assertIsNotNone(usd_context.get_stage())

        vp_window = ui.Window('test_capture_file_with_format', width=TEST_WIDTH, height=TEST_HEIGHT, flags=ui.WINDOW_FLAGS_NO_SCROLLBAR)
        self.assertIsNotNone(vp_window)
        with vp_window.frame:
            vp_widget = ViewportWidget(resolution=(TEST_WIDTH, TEST_HEIGHT))
            self.assertIsNotNone(vp_widget)

            format_desc = {}
            format_desc["format"] = "exr"
            format_desc["compression"] = "rle"

            capture = vp_widget.viewport_api.schedule_capture(FileCapture(image_path, format_desc=format_desc))
            captured_aovs = await capture.wait_for_result()
            self.assertEqual(captured_aovs, ['LdrColor'])
            # Test the frame's metadata is available
            self.assertIsNotNone(capture.view)
            self.assertIsNotNone(capture.projection)
            self.assertEqual(capture.resolution, (TEST_WIDTH, TEST_HEIGHT))
            # Test the viewport_handle is also available for usets that
            self.assertIsNotNone(capture.viewport_handle)
            self.assertIsNotNone(capture.format_desc)

            # Cannot compare EXR images for testing yet
            # compare_images(image_name)

        vp_widget.destroy()
        vp_window.destroy()
        del vp_window

    async def merged_test_capture_bytes_free_func_hdr(self):
        # 1030: Doesn't have the ability to set render_product_path and this test will fail
        import omni.hydratexture
        if not hasattr(omni.hydratexture.IHydraTexture, 'set_render_product_path'):
            return

        await self.create_test_area(width=TEST_WIDTH, height=TEST_HEIGHT)

        aov_name = 'HdrColor'
        image_name = f'cube_{aov_name}.exr'
        image_path = str(OUTPUTS_DIR.joinpath(image_name))
        usd_path = str(TEST_FILES_DIR.joinpath('usd/cube.usda'))

        usd_context = omni.usd.get_context()
        await usd_context.open_stage_async(usd_path)
        self.assertIsNotNone(usd_context.get_stage())

        num_calls = 0
        def on_capture_completed(buffer, buffer_size, width, height, format):
            nonlocal num_calls
            num_calls = num_calls + 1
            self.assertEqual(width, TEST_WIDTH)
            self.assertEqual(height, TEST_HEIGHT)
            # Note the RGBA16_SFLOAT format!
            self.assertEqual(format, ui.TextureFormat.RGBA16_SFLOAT)

        vp_window = ui.Window('test_capture_bytes_free_func', width=TEST_WIDTH, height=TEST_HEIGHT, flags=ui.WINDOW_FLAGS_NO_SCROLLBAR)
        self.assertIsNotNone(vp_window)
        with vp_window.frame:
            vp_widget = ViewportWidget(resolution=(TEST_WIDTH, TEST_HEIGHT))
            self.assertIsNotNone(vp_widget)

            # Edit the RenderProduct to output HdrColor
            render_product_path = vp_widget.viewport_api.render_product_path
            if render_product_path[0] != '/':
                render_product_path = f'/Render/RenderProduct_{render_product_path}'

            stage = usd_context.get_stage()
            with Usd.EditContext(stage, stage.GetSessionLayer()):
                prim = usd_context.get_stage().GetPrimAtPath(render_product_path)
                self.assertTrue(prim.IsValid())
                self.assertTrue(prim.IsA(UsdRender.Product))
                product = UsdRender.Product(prim)
                targets = product.GetOrderedVarsRel().GetForwardedTargets()
                self.assertEqual(1, len(targets))

                prim = usd_context.get_stage().GetPrimAtPath(targets[0])
                self.assertTrue(prim.IsValid())
                self.assertTrue(prim.IsA(UsdRender.Var))
                render_var = UsdRender.Var(prim)
                value_set = render_var.GetSourceNameAttr().Set(aov_name)
                self.assertTrue(value_set)

            # Trigger the texture to update the render, and wait for the change to funnel back to us
            vp_widget.viewport_api.render_product_path = render_product_path
            settings_changed = await vp_widget.viewport_api.wait_for_render_settings_change()
            self.assertTrue(settings_changed)

            # Now that the settings have changed, so the capture of 'HdrColor'
            captured_aovs = await vp_widget.viewport_api.schedule_capture(ByteCapture(on_capture_completed, aov_name=aov_name)).wait_for_result()
            self.assertTrue(aov_name in captured_aovs)

            # Since we requested one AOV, test that on_capture_completed was only called once
            # Currently RTX will deliver 'HdrColor' and 'LdrColor' when 'HdrColor' is requested
            # If that changes these two don't neccessarily need to be tested.
            self.assertEqual(len(captured_aovs), 1)
            self.assertEqual(num_calls, 1)

            captured_aovs = await vp_widget.viewport_api.schedule_capture(FileCapture(image_path, aov_name=aov_name)).wait_for_result()
            self.assertTrue(aov_name in captured_aovs)

            # XXX cannot compare EXR
            # compare_images(image_name)
            # So test it exists and publish the image
            self.assertTrue(os.path.exists(image_path))
            teamcity_publish_image_artifact(image_path, "results")

        vp_widget.destroy()
        vp_window.destroy()
        del vp_window

    async def merged_test_capture_bytes_subclass(self):
        await self.create_test_area(width=TEST_WIDTH, height=TEST_HEIGHT)

        usd_path = str(TEST_FILES_DIR.joinpath('usd/cube.usda'))
        usd_context = omni.usd.get_context()
        await usd_context.open_stage_async(usd_path)
        self.assertIsNotNone(usd_context.get_stage())

        capture = TestByteCapture(self)

        vp_window = ui.Window('test_capture_bytes_subclass', width=TEST_WIDTH, height=TEST_HEIGHT, flags=ui.WINDOW_FLAGS_NO_SCROLLBAR)
        self.assertIsNotNone(vp_window)
        with vp_window.frame:
            vp_widget = ViewportWidget(resolution=(TEST_WIDTH, TEST_HEIGHT))
            self.assertIsNotNone(vp_widget)
            vp_widget.viewport_api.schedule_capture(capture)

        captured_aovs = await capture.wait_for_result()
        self.assertEqual(captured_aovs, ['LdrColor'])

        vp_widget.destroy()
        vp_window.destroy()
        del vp_window

    async def test_3_capture_tests_merged_for_linux(self):
        try:
            await self.merged_test_capture_file()
            await self.wait_n_updates(10)

            await self.merged_test_capture_bytes_free_func_hdr()
            await self.wait_n_updates(10)

            await self.merged_test_capture_bytes_subclass()
            await self.wait_n_updates(10)

            await self.merged_test_capture_file_with_format()
            await self.wait_n_updates(10)
        finally:
            await self.finalize_test_no_image()
