# 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 carb
import carb.input
import omni.appwindow
import omni.kit.commands
import omni.kit.test
import omni.kit.viewport_legacy
import omni.usd
from carb.input import KeyboardInput as Key
from omni.kit.test_suite.helpers import arrange_windows


class TestViewportMouseRayCast(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows(hide_viewport=True)
        self._input_provider = carb.input.acquire_input_provider()
        app_window = omni.appwindow.get_default_app_window()
        self._mouse = app_window.get_mouse()
        self._keyboard = app_window.get_keyboard()
        self._app = omni.kit.app.get_app()

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

        viewport_interface = omni.kit.viewport_legacy.acquire_viewport_interface()
        self._viewport = viewport_interface.get_viewport_window()
        self._viewport.set_enabled_picking(True)
        self._viewport.set_window_pos(0, 0)
        self._viewport.set_window_size(400, 400)
        self._viewport.show_hide_window(True)

        await self._app.next_update_async()

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

    async def test_viewport_mouse_raycast(self):
        stage_update = omni.stageupdate.get_stage_update_interface()
        stage_subscription = stage_update.create_stage_update_node(
            "viewport_mouse_raycast_test", on_raycast_fn=self._on_ray_cast
        )

        test_data = [
            (
                "/OmniverseKit_Persp",
                (499.058, 499.735, 499.475),
                (-0.848083, -0.23903, -0.472884),
            ),
            (
                "/OmniverseKit_Front",
                (-127.551, 165.44, 30000),
                (0, 0, -1),
            ),
            (
                "/OmniverseKit_Top",
                (127.551, 30000, 165.44),
                (0, -1, 0),
            ),
            (
                "/OmniverseKit_Right",
                (-30000, 165.44, -127.551),
                (1, 0, 0),
            ),
        ]

        for data in test_data:
            self._raycast_data = None

            self._viewport.set_active_camera(data[0])
            await self._app.next_update_async()

            self._input_provider.buffer_mouse_event(
                self._mouse, carb.input.MouseEventType.MOVE, (100, 100), 0, (100, 100)
            )
            await self._app.next_update_async()

            # Hold SHIFT down to start raycast
            self._input_provider.buffer_keyboard_key_event(
                self._keyboard, carb.input.KeyboardEventType.KEY_PRESS, Key.LEFT_SHIFT, 0
            )
            await self._app.next_update_async()

            self._input_provider.buffer_keyboard_key_event(
                self._keyboard, carb.input.KeyboardEventType.KEY_RELEASE, Key.LEFT_SHIFT, 0
            )
            await self._app.next_update_async()

            self._input_provider.buffer_mouse_event(self._mouse, carb.input.MouseEventType.MOVE, (0, 0), 0, (0, 0))
            await self._app.next_update_async()

            self.assertIsNotNone(self._raycast_data)

            self.assertAlmostEqual(self._raycast_data[0][0], data[1][0], 3)
            self.assertAlmostEqual(self._raycast_data[0][1], data[1][1], 3)
            self.assertAlmostEqual(self._raycast_data[0][2], data[1][2], 3)

            self.assertAlmostEqual(self._raycast_data[1][0], data[2][0], 3)
            self.assertAlmostEqual(self._raycast_data[1][1], data[2][1], 3)
            self.assertAlmostEqual(self._raycast_data[1][2], data[2][2], 3)

        stage_subscription = None

    def _on_ray_cast(self, orig, dir, input):
        self._raycast_data = (orig, dir)
        print(self._raycast_data)
