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

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 AsyncRenderingTest(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):
        print("Tearing down..")
        self._usd_context.close_stage()
        omni.usd.release_all_hydra_engines(self._usd_context)

        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 test_1_simple_context_attach(self):
        renderer = "rtx"
        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("simple_cubes_mat.usda")
        print("Opening '%s'" % (test_usd_asset))

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

        app = omni.kit.app.get_app()
        self._update_counter = 0
        self._render_counter = 0
        self._asset_loaded = False
        def on_update(event):
            self._update_counter += 1
        self._app_update_sub = app.get_update_event_stream().create_subscription_to_pop(on_update, name='async rendering test update')

        is_async = self._settings.get("/app/asyncRendering")
        print("Async is %s" % (is_async))
        self._hydra_texture = self._hydra_texture_factory.create_hydra_texture(
            "test_viewport",
            1280,
            720,
            self._usd_context_name,
            "/test_cam",
            renderer,
            is_async=is_async
        )

        def _on_stage_event(event):
            ASSETS_LOADED = int(omni.usd.StageEventType.ASSETS_LOADED)
            if (self._asset_loaded == False) and (event.type is ASSETS_LOADED):
                self._asset_loaded = True
                print("Assets loaded. UI frame: %d" % (self._update_counter))

        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(_on_stage_event, name="stage events")

        def on_drawable_changed(event: carb.events.IEvent):
            # Renderer counter should start after usd assets are completely loaded in async mode
            if self._asset_loaded or (is_async == False):
                # +1 because the frame number is zero-based.
                self._render_counter = self._hydra_texture.get_frame_info(event.payload['result_handle']).get('frame_number') + 1
                print("Rendered %d frames (UI frames drawn: %d)" % (self._render_counter, self._update_counter))

        self._drawable_changed_sub = self._hydra_texture.get_event_stream().create_subscription_to_push_by_type(
            omni.hydratexture.EVENT_TYPE_DRAWABLE_CHANGED,
            on_drawable_changed,
            name='async rendering test drawable update',
        )

        print("Waiting the RTX renderer to load and start draw pictures")
        MAX_WAIT_FRAMES_INIT = 50000
        wait_frames_left = MAX_WAIT_FRAMES_INIT
        while True:
            await app.next_update_async()
            if self._render_counter > 0 or wait_frames_left <= 0:
                break
            wait_frames_left -= 1
        print("Waited %d frames before renderer initialized" % (MAX_WAIT_FRAMES_INIT - wait_frames_left))

        # Set render mode and settings
        self._settings.set("/rtx/rendermode", "PathTracing")
        self._settings.set("/rtx/pathtracing/totalSpp", 0)
        self._settings.set("/rtx/pathtracing/spp", 256)
        self._settings.set("/rtx/pathtracing/clampSpp", 256)

        print("Starting render loop")
        test_updates_before = self._update_counter
        test_renders_before = self._render_counter

        wait_frames_left = 10000
        while True:
            await app.next_update_async()
            if self._render_counter - test_renders_before > 10 or wait_frames_left <= 0:
                break
            wait_frames_left -= 1

        self.assertTrue(self._render_counter > 10)

        test_updates_delta = self._update_counter - test_updates_before
        test_renders_delta = self._render_counter - test_renders_before

        updates_renders_ratio = test_updates_delta / test_renders_delta
        updates_renders_difference = test_updates_delta - test_renders_delta
        print("Updates: %d, renders: %d, difference %d, ratio %3.2f" % (test_updates_delta, test_renders_delta, updates_renders_difference, updates_renders_ratio))

        # Verify results, this depends on the Async mode
        if is_async:
            # When Async is On - we expect much more UI updates than render delegate renders
            # R525+ drivers showing 3X+ UI draws than 10X+ in comparison to R470.
            # OM-78715: lowering threshold even further since it occasionally fails on IPP agents.
            self.assertTrue(updates_renders_ratio > 2.3)
        else:
            # When Async is Off - we expect UI updates come in a lockstep with render delegate renders
            # Add some leeway for potential frame offsetting
            self.assertTrue(updates_renders_ratio < 1.9)

            NUM_FRAMES_IN_FLIGHT = 3
            # The UI should be no more than NUM_FRAMES_IN_FLIGHT + 1 frames ahead of the renderer.
            self.assertTrue(0 <= updates_renders_difference and updates_renders_difference <= NUM_FRAMES_IN_FLIGHT + 1)

        self._drawable_changed_sub = None
        self._hydra_texture = None
        self._app_update_sub = None
        self._stage_event_sub = None
