## Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
##
## NVIDIA CORPORATION and its licensors retain all intellectual property
## and proprietary rights in and to this software, related documentation
## and any modifications thereto.  Any use, reproduction, disclosure or
## distribution of this software and related documentation without an express
## license agreement from NVIDIA CORPORATION is strictly prohibited.
##
import omni.kit.app
import omni.kit.commands
import omni.kit.undo
import omni.kit.test
import omni.usd
import carb
from .test_common import RtxTest, testSettings, postLoadTestSettings, wait_for_update

class TestRtxPicking(RtxTest):
    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        omni.usd.get_context().new_stage()  # Important: it should be called before `setUp` method as it reset settings and stage.
        await omni.kit.app.get_app().next_update_async()  # Wait stage loading
        self.set_settings(postLoadTestSettings)

    async def test_picking_queue(self):
        """
        Verify picking/query requests are successfully queued
        """
        super().add_floor()
        await wait_for_update()

        # OM-80470: Verify all picking requests requested in consecutive frames are all queued and executed successfully
        counter = 0
        def query_complete(path, pos, *args, **kwargs):
            nonlocal counter
            counter += 1

        import omni.kit.viewport
        viewport = omni.kit.viewport.utility.get_active_viewport()

        if hasattr(viewport, 'legacy_window'):
            return

        # Queue multiple query requests interrupted by a pick
        center = [int(v * 0.5) for v in  viewport.resolution]
        viewport.request_query(center, query_complete, query_name="test_1")
        viewport.request_query(center, query_complete, query_name="test_2")
        viewport.request_pick(center, center, omni.usd.PickingMode.RESET_AND_SELECT)
        viewport.request_query(center, query_complete, query_name="test_3")
        viewport.request_query(center, query_complete, query_name="test_4")

        await wait_for_update(wait_frames=40)

        # All queries should complete
        self.assertEqual(counter, 4)
