import os
import sys
import inspect
import importlib

import carb
import carb.settings
import carb.tokens

import omni.kit.app
import omni.kit.test
#import omni.kit.test_helpers_gfx

import omni.kit.renderer.bind

from omni.kit.test.teamcity import teamcity_publish_image_artifact

OUTPUTS_DIR = omni.kit.test.get_test_output_path()
USE_TUPLES = True

class AppDevCompatTest(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._settings = carb.settings.acquire_settings_interface()
        self._app_window_factory = omni.appwindow.acquire_app_window_factory_interface()
        self._renderer = omni.kit.renderer.bind.acquire_renderer_interface()

        self._module_test_helpers_gfx = importlib.import_module('omni.kit.test_helpers_gfx')

        self._captured_buffer = []
        self._captured_buffer_w = 0
        self._captured_buffer_h = 0
        self._captured_buffer_fmt = 0

    def __test_name(self) -> str:
        return f"{self.__module__}.{self.__class__.__name__}.{inspect.stack()[2][3]}"

    async def tearDown(self):
        self._renderer = None
        self._app_window_factory = None
        self._settings = None

    def _capture_callback(self, buf, buf_size, w, h, fmt):
        if USE_TUPLES:
            self._captured_buffer = omni.renderer_capture.convert_raw_bytes_to_rgba_tuples(buf, buf_size, w, h, fmt)
        else:
            self._captured_buffer = omni.renderer_capture.convert_raw_bytes_to_list(buf, buf_size, w, h, fmt)
        self._captured_buffer_w = w
        self._captured_buffer_h = h
        self._captured_buffer_fmt = fmt

    async def test_1_capture_variance(self):
        test_name = self.__test_name()

        import omni.renderer_capture

        app_window = self._app_window_factory.get_default_window()

        capture_interface = omni.renderer_capture.acquire_renderer_capture_interface()
        capture_interface.capture_next_frame_swapchain_callback(self._capture_callback, app_window)

        await omni.kit.app.get_app().next_update_async()

        capture_interface.wait_async_capture(app_window)

        if "PIL" not in sys.modules.keys():
            # Checking if we have Pillow imported
            try:
                from PIL import Image
            except ImportError:
                # Install Pillow if it's not installed
                import omni.kit.pipapi

                omni.kit.pipapi.install("Pillow", module="PIL")

        from PIL import Image
        from PIL import ImageStat

        image = Image.new('RGBA', [self._captured_buffer_w, self._captured_buffer_h])

        if USE_TUPLES:
            image.putdata(self._captured_buffer)
        else:
            buf_channel_it = iter(self._captured_buffer)
            captured_buffer_tuples = list(zip(buf_channel_it, buf_channel_it, buf_channel_it, buf_channel_it))
            image.putdata(captured_buffer_tuples)

        # Only compare the standard deviation, since the test failing because of the UI elements shifting around is
        # no good. Generally, it is enough to test that:
        #   1. the test doesn't crash on iGPUs,
        #   2. the test produces anything other than the black screen (the std dev of which is 0.0)
        image_stats = ImageStat.Stat(image)
        avg_std_dev = sum(image_stats.stddev) / 3.0

        # Save image for user verification
        test_name = self.__test_name()
        TEST_IMG_PATH = os.path.join(OUTPUTS_DIR, test_name + ".png")
        image.save(TEST_IMG_PATH)

        # The full app.dev has a std dev ~30.0, and the version with broken layout has a std dev of ~18.0
        STD_DEV_THRESHOLD = 10.0
        if avg_std_dev < STD_DEV_THRESHOLD:
            teamcity_publish_image_artifact(TEST_IMG_PATH, "results", "Generated")
            carb.log_error("Standard deviation of the produced image is lower than the threshold!")
            self.assertTrue(False)
