import pathlib
import asyncio

import carb
import carb.settings
import carb.tokens

import omni.kit.app
import omni.kit.test

import omni.hydratexture
import omni.usd
from typing import Callable

# FIXME: omni.ui.ImageProvider holds the carb.Format conversion routine
import omni.ui
from pxr import Gf, UsdRender

EXTENSION_FOLDER_PATH = pathlib.Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))
DATA_DIR = EXTENSION_FOLDER_PATH.joinpath("data/tests")


class HydraTextureTest(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._settings = carb.settings.acquire_settings_interface()
        self._hydra_texture_factory = omni.hydratexture.acquire_hydra_texture_factory_interface()

        self._usd_context_name = ''
        self._usd_context = omni.usd.get_context(self._usd_context_name)
        await self._usd_context.new_stage_async()

    async def tearDown(self):
        self._hydra_texture_factory = None
        self._settings = None

        wait_iterations = 6
        for i in range(wait_iterations):
            await omni.kit.app.get_app().next_update_async()

    async def _create_hydra_texture_test(self, filename: str, texture_test: Callable,
                                         renderer: str = 'pxr', res_x: int = 320, res_y: int = 320,
                                         engine_options: dict = None):
        wait_iterations = 6
        try:
            if renderer not in self._usd_context.get_attached_hydra_engine_names():
                omni.usd.add_hydra_engine(renderer, self._usd_context)

            test_usd_asset = DATA_DIR.joinpath(filename)
            print("Opening '%s'" % (test_usd_asset))

            await self._usd_context.open_stage_async(str(test_usd_asset))

            if engine_options is None:
                engine_options = {"is_async": False}

            hydra_texture = self._hydra_texture_factory.create_hydra_texture(
                "test_viewport",
                res_x,
                res_y,
                self._usd_context_name,
                "/test_cam",
                renderer,
                **engine_options
            )

            return await texture_test(hydra_texture)
        finally:
            for i in range(wait_iterations):
                await omni.kit.app.get_app().next_update_async()

    # test1 - test1g simulate possible failure on low memory Linux configs from
    # bad gpu-foundation shutdown and startup.
    async def test_1_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1b_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1c_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1d_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1e_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1f_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_1g_simple_context_attach(self):
        async def no_more_test(hydra_texture):
            pass
        await self._create_hydra_texture_test('simple_cubes_mat.usda', no_more_test)

    async def test_2_drawable_changed(self):
        drawable_result = asyncio.Future()
        async def drawable_changed_test(hydra_texture):
            def on_drawable_changed(event: carb.events.IEvent):
                if event.type != omni.hydratexture.EVENT_TYPE_DRAWABLE_CHANGED:
                    carb.log_error("Wrong event captured for DRAWABLE_CHANGED!")
                    return

                result_handle = event.payload['result_handle']
                aov_info = hydra_texture.get_aov_info(result_handle)
                self.assertEqual(aov_info[0]['name'], 'LdrColor')

                ldr_info = hydra_texture.get_aov_info(result_handle, 'LdrColor')
                self.assertEqual(ldr_info[0]['name'], 'LdrColor')

                ldr_info = hydra_texture.get_aov_info(result_handle, 'LdrColor', include_texture=True)
                self.assertEqual(ldr_info[0]['name'], 'LdrColor')
                ldr_tex = ldr_info[0]['texture']
                self.assertEqual(ldr_tex['resolution'][0], 640)
                self.assertEqual(ldr_tex['resolution'][1], 320)
                self.assertEqual(ldr_tex['format'], omni.ui.TextureFormat.RGBA8_UNORM)
                self.assertIsNotNone(ldr_tex['rp_resource'])

                frame_info = hydra_texture.get_frame_info(result_handle)
                self.assertIsNotNone(frame_info['view'])
                self.assertIsNotNone(frame_info['projection'])
                self.assertIsNotNone(frame_info['fps'])
                self.assertIsNotNone(frame_info['resolution'])
                self.assertIsNotNone(frame_info['progression'])
                self.assertIsNotNone(frame_info['frame_number'])
                self.assertIsNotNone(frame_info['device_mask'])

                nonlocal drawable_result
                # EVENT_TYPE_DRAWABLE_CHANGED can be called multiple times before yielding for async loop
                if not drawable_result.done():
                    drawable_result.set_result(True)

            return hydra_texture.get_event_stream().create_subscription_to_push_by_type(
                omni.hydratexture.EVENT_TYPE_DRAWABLE_CHANGED,
                on_drawable_changed,
                name="Viewport Texture drawable change",
            )

        drawable_change_sub = await self._create_hydra_texture_test('simple_cubes_mat.usda', drawable_changed_test,
                                                                    res_x = 640, res_y = 320)
        result = await drawable_result
        drawable_change_sub = None
        self.assertTrue(result)

    async def test_3_custom_product(self):
        # Build up the custom RenderProduct
        custom_prod_path = '/Render/renderproduct_custom'
        custom_var_path = '/Render/rendervar_custom'
        custom_cam_path = '/test_cam_custom'
        custom_resolution = Gf.Vec2i(512, 512)

        # XXX: This used to be handled entirely by set_render_product_path and usd-abi
        # But usd-abi now validates portions of the RenderProduct.
        # So just set up the complete UsdRender.Product and UsdRender.Var and validate it ourselves along the way.
        async def setup_custom_product(stage):
            custom_prod = UsdRender.Product.Define(stage, custom_prod_path)
            self.assertIsNotNone(custom_prod)

            custom_var = UsdRender.Var.Define(stage, custom_var_path)
            self.assertIsNotNone(custom_var)

            ordered_vars = custom_prod.GetOrderedVarsRel()
            self.assertIsNotNone(ordered_vars)
            ordered_vars.SetTargets([custom_var_path])
            self.assertEqual([sdf_path.pathString for sdf_path in ordered_vars.GetForwardedTargets()], [custom_var_path])

            camera_rel = custom_prod.GetCameraRel()
            self.assertIsNotNone(camera_rel)
            camera_rel.SetTargets([custom_cam_path])
            self.assertEqual([sdf_path.pathString for sdf_path in camera_rel.GetForwardedTargets()], [custom_cam_path])

            res_attr = custom_prod.GetResolutionAttr()
            self.assertIsNotNone(res_attr)
            res_attr.Set(custom_resolution)

            source_name = custom_var.GetSourceNameAttr()
            self.assertIsNotNone(source_name)
            source_name.Set('LdrColor')

            # Wait one frame to make sure Usd edits are absorbed
            return await omni.kit.app.get_app().next_update_async()

        async def product_test(hydra_texture):
            dflt_path = hydra_texture.get_render_product_path()

            await setup_custom_product(self._usd_context.get_stage())

            # Move the product to a custom one keeping width, height, and camera
            success = hydra_texture.set_render_product_path(custom_prod_path, keep_camera=True, keep_resolution=True)
            cstm_path = hydra_texture.get_render_product_path()
            self.assertTrue(success)
            self.assertNotEqual(dflt_path, cstm_path)

            self.assertEqual(hydra_texture.get_camera_path(), '/test_cam')
            self.assertEqual(hydra_texture.get_width(), 320)
            self.assertEqual(hydra_texture.get_height(), 320)

            # Move the product to a custom one, reading width, height, and camera
            success = hydra_texture.set_render_product_path(custom_prod_path)
            cstm_path = hydra_texture.get_render_product_path()
            self.assertTrue(success)
            self.assertNotEqual(dflt_path, cstm_path)

            # XXX: Can't test this here, needs to process the Viewport
            # self.assertEqual(hydra_texture.get_camera_path(), custom_cam_path)
            self.assertEqual(hydra_texture.get_width(), custom_resolution[0])
            self.assertEqual(hydra_texture.get_height(), custom_resolution[1])

            # Move it back to the default
            success = hydra_texture.set_render_product_path(dflt_path)
            self.assertTrue(success)
            self.assertEqual(dflt_path, hydra_texture.get_render_product_path())

        await self._create_hydra_texture_test('custom_product.usda', product_test)

    async def test_tick_rate_creation(self):
        """Test arguments passed to create_hydra_texture"""
        engine_options = {
            "is_async": False,
            "hydra_tick_rate": 30
        }
        async def check_engine_tick_rate(hydra_texture):
            settings = carb.settings.get_settings()
            tick_rate = settings.get(f"{hydra_texture.get_settings_path()}hydraTickRate")
            is_async = settings.get(f"{hydra_texture.get_settings_path()}async")

            self.assertEqual(tick_rate, engine_options.get("hydra_tick_rate"))
            self.assertEqual(is_async, engine_options.get("is_async"))

        await self._create_hydra_texture_test('simple_cubes_mat.usda', check_engine_tick_rate, engine_options=engine_options)

    async def test_async_engine_creation(self):
        """Test arguments passed to create_hydra_texture"""
        engine_options = {
            "is_async": True,
        }
        async def check_engine_async(hydra_texture):
            settings = carb.settings.get_settings()
            is_async = settings.get(f"{hydra_texture.get_settings_path()}async")
            self.assertEqual(is_async, engine_options.get("is_async"))

        await self._create_hydra_texture_test('simple_cubes_mat.usda', check_engine_async, engine_options=engine_options)

    # There isn't a great mechanism for having a test fail mean success; which whould allow
    # these two tests that a failure wont result in a crash.
    # Until then, hopefully these can be run on new test additions or manually.
    #
    # async def test_0_disabled_test_failure_wont_crash(self):
    #     self.assertTrue(False)
    # async def test_z_disabled_test_failure_wont_crash(self):
    #     self.assertTrue(False)
