import asyncio
import carb
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, wait_stage_loading


# 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 ViewportContextTest(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()
        self._usd_context = omni.usd.get_context()
        await self._usd_context.new_stage_async()

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

    async def _create_test_stage(self):
        # Create the objects
        test_list = omni.kit.viewport_legacy.tests.create_test_stage()

        # Wait so that the test does not exit before MDL loaded and crash on exit.
        while True:
            try:
                event, payload = await asyncio.wait_for(self._usd_context.next_stage_event_async(), timeout=60.0)
                if event == int(omni.usd.StageEventType.ASSETS_LOADED):
                    break
            except asyncio.TimeoutError:
                _, files_loaded, total_files = self._usd_context.get_stage_loading_status()
                if files_loaded == total_files:
                    carb.log_warn("Timed out waiting for ASSETS_LOADED event. Is MDL already loaded?")
                    break

        return test_list

    # Simulate a mouse-click-selection at pos (x,y) in viewport
    async def _simulate_mouse_selection(self, viewport, pos):
        app = omni.kit.app.get_app()
        mouse = omni.appwindow.get_default_app_window().get_mouse()
        input_provider = carb.input.acquire_input_provider()

        async def simulate_mouse_click():
            input_provider.buffer_mouse_event(mouse, carb.input.MouseEventType.MOVE, pos, 0, pos)
            input_provider.buffer_mouse_event(mouse, carb.input.MouseEventType.LEFT_BUTTON_DOWN, pos, 0, pos)
            await app.next_update_async()

            input_provider.buffer_mouse_event(mouse, carb.input.MouseEventType.LEFT_BUTTON_UP, pos, 0, pos)
            await app.next_update_async()

        # You're guess as good as mine.
        # Need to start and end with one wait, then need 12 more waits for mouse click
        loop_counts = (1, 5)

        for i in range(loop_counts[0]):
            await app.next_update_async()

        # Is this really to sync up carb and imgui state ?
        for i in range(loop_counts[1]):
            await simulate_mouse_click()

        # Enable picking and do the selection
        viewport.request_picking()
        await simulate_mouse_click()

        # Almost done!
        for i in range(loop_counts[0]):
            await app.next_update_async()

    # Actual test, notice it is "async" function, so "await" can be used if needed
    #
    # But why does this test exist in Viewport??
    async def test_bound_objects(self):
        return
        # Create the objects
        usd_context, test_list = self._usd_context, await self._create_test_stage()
        self.assertIsNotNone(usd_context)
        self.assertIsNotNone(test_list)

        objects = {}
        stage = omni.usd.get_context().get_stage()
        objects["stage"] = stage
        for mtl, mtl_prims in test_list:

            # Select material prim
            objects["prim_list"] = [stage.GetPrimAtPath(Sdf.Path(mtl))]

            # Run select_prims_using_material
            omni.kit.context_menu.get_instance().select_prims_using_material(objects)

            # Verify selected prims
            prims = omni.usd.get_context().get_selection().get_selected_prim_paths()
            self.assertTrue(prims == mtl_prims)

    async def test_create_and_select(self):
        viewport_interface = omni.kit.viewport_legacy.acquire_viewport_interface()
        viewport = viewport_interface.get_viewport_window()
        viewport.set_enabled_picking(True)
        viewport.set_window_pos(0, 0)
        viewport.set_window_size(400, 400)

        # Create the objects
        usd_context, cube_path = self._usd_context, omni.kit.viewport_legacy.tests.create_and_select_cube()
        self.assertIsNotNone(usd_context)
        self.assertIsNotNone(cube_path)

        self.assertEqual([], self._usd_context.get_selection().get_selected_prim_paths())

        await self._simulate_mouse_selection(viewport, (200, 230))

        self.assertEqual([cube_path], self._usd_context.get_selection().get_selected_prim_paths())
