import carb.dictionary
import omni.usd
import omni.kit.test
import omni.kit.viewport_legacy
import omni.timeline
from pxr import Usd, UsdGeom, Gf, Sdf
from omni.kit.test_suite.helpers import arrange_windows

# Having a test class dervived from omni.kit.test.AsyncTestCase declared on the root of module will make it auto-discoverable by omni.kit.test
class ViewportDrawTest(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()

    # After running each test
    async def tearDown(self):
        pass

    # Actual test, notice it is "async" function, so "await" can be used if needed
    async def test_ui_draw_event_payload(self):
        dictionary_interface = carb.dictionary.get_dictionary()
        viewport_interface = omni.kit.viewport_legacy.acquire_viewport_interface()
        viewport = viewport_interface.get_viewport_window()
        draw_event_stream = viewport.get_ui_draw_event_stream()

        CAM_PATH = "/OmniverseKit_Persp"
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        cam_prim = stage.GetPrimAtPath(CAM_PATH)
        camera = UsdGeom.Camera(cam_prim)

        timeline_interface = omni.timeline.get_timeline_interface()
        timecode = timeline_interface.get_current_time() * stage.GetTimeCodesPerSecond()

        gf_camera = camera.GetCamera(timecode)

        expected_transform = gf_camera.transform
        expected_viewport_rect = viewport.get_viewport_rect()

        def uiDrawCallback(event):
            events_data = dictionary_interface.get_dict_copy(event.payload)
            vm = events_data["viewMatrix"]
            viewMatrix = Gf.Matrix4d(
                vm[0],
                vm[1],
                vm[2],
                vm[3],
                vm[4],
                vm[5],
                vm[6],
                vm[7],
                vm[8],
                vm[9],
                vm[10],
                vm[11],
                vm[12],
                vm[13],
                vm[14],
                vm[15],
            ).GetInverse()
            self.assertTrue(Gf.IsClose(expected_transform, viewMatrix, 0.00001))

            vr = events_data["viewportRect"]
            self.assertAlmostEqual(expected_viewport_rect[0], vr[0])
            self.assertAlmostEqual(expected_viewport_rect[1], vr[1])
            self.assertAlmostEqual(expected_viewport_rect[2], vr[2])
            self.assertAlmostEqual(expected_viewport_rect[3], vr[3])

        subscription = draw_event_stream.create_subscription_to_pop(uiDrawCallback)
        await omni.kit.app.get_app().next_update_async()
