## Copyright (c) 2022, 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.test
import omni.ui as ui
import omni.kit.undo
from omni.ui.tests.test_base import OmniUiTest
from pxr import Kind, Sdf, Gf
from omni.kit import ui_test
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, wait_stage_loading, select_prims, arrange_windows


class TestBoolType(OmniUiTest):
    def __init__(self, tests=()):
        super().__init__(tests)

    # Before running each test
    async def setUp(self):
        await arrange_windows(topleft_window="Property", topleft_height=256, topleft_width=800.0)
        await open_stage(get_test_data_path(__name__, "usd/bools.usda"))
        await self._show_raw(False)

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

    async def _show_raw(self, new_state):
        await select_prims(["/defaultPrim/in_0"])
        for w in ui_test.find_all("Property//Frame/**/CollapsableFrame[*]"):
            if w.widget.title == "Raw USD Properties":
                w.widget.collapsed = new_state
        # change prim selection immediately after collapsed state change can result in one bad frame
        # i.e. selection change event happens before frame UI rebuild, and old widget models are already destroyed.
        await ui_test.human_delay()
        await select_prims([])

    async def test_l1_bool_test(self):
        # setup
        await wait_stage_loading()
        stage = omni.usd.get_context().get_stage()
        prim = stage.GetPrimAtPath('/World/Looks/bool_types/Shader')
        await select_prims(["/World/Looks/bool_types/Shader"])
        await ui_test.human_delay()

        # test bool checkbox
        attr = prim.GetAttribute("inputs:MyBool")
        self.assertTrue(attr.Get())
        await ui_test.find("Property//Frame/**/.identifier=='bool_inputs:MyBool'").click()
        self.assertFalse(attr.Get())
        await ui_test.find("Property//Frame/**/.identifier=='control_state_inputs:MyBool'").click()
        self.assertTrue(attr.Get())

        # test bool vec2 checkboxes
        attr = prim.GetAttribute("inputs:MyBool2")
        self.assertEqual(attr.Get(), (0, 0))
        widget = ui_test.find("Property//Frame/**/.identifier=='boolean_per_channel_inputs:MyBool2'")
        for index, expected in enumerate([(0, 0), (0, 1), (1, 0), (1, 1)]):
            checkwidgets = widget.find_all("**/CheckBox[*]")
            if index & 1:
                await checkwidgets[1].click()
            if index & 2:
                await checkwidgets[0].click()

            self.assertEqual(attr.Get(), expected)
            await ui_test.find("Property//Frame/**/.identifier=='control_state_inputs:MyBool2'").click()
            self.assertEqual(attr.Get(), (0, 0))
            omni.kit.undo.undo()
            self.assertEqual(attr.Get(), expected)
            omni.kit.undo.redo()
            self.assertEqual(attr.Get(), (0, 0))

        # test bool vec3 checkboxes
        attr = prim.GetAttribute("inputs:MyBool3")
        self.assertEqual(attr.Get(), (0, 0, 0))
        widget = ui_test.find("Property//Frame/**/.identifier=='boolean_per_channel_inputs:MyBool3'")
        for index, expected in enumerate([(0, 0, 0),
                                          (0, 0, 1),
                                          (0, 1, 0),
                                          (0, 1, 1),
                                          (1, 0, 0),
                                          (1, 0, 1),
                                          (1, 1, 0),
                                          (1, 1, 1)]):
            checkwidgets = widget.find_all("**/CheckBox[*]")
            if index & 1:
                await checkwidgets[2].click()
            if index & 2:
                await checkwidgets[1].click()
            if index & 4:
                await checkwidgets[0].click()

            self.assertEqual(attr.Get(), expected)
            await ui_test.find("Property//Frame/**/.identifier=='control_state_inputs:MyBool3'").click()
            self.assertEqual(attr.Get(), (0, 0, 0))
            omni.kit.undo.undo()
            self.assertEqual(attr.Get(), expected)
            omni.kit.undo.redo()
            self.assertEqual(attr.Get(), (0, 0, 0))

        # test bool vec4 checkboxes
        attr = prim.GetAttribute("inputs:MyBool4")
        self.assertEqual(attr.Get(), (0, 0, 0, 0))
        widget = ui_test.find("Property//Frame/**/.identifier=='boolean_per_channel_inputs:MyBool4'")
        for index, expected in enumerate([(0, 0, 0, 0),
                                          (0, 0, 0, 1),
                                          (0, 0, 1, 0),
                                          (0, 0, 1, 1),
                                          (0, 1, 0, 0),
                                          (0, 1, 0, 1),
                                          (0, 1, 1, 0),
                                          (0, 1, 1, 1),
                                          (1, 0, 0, 0),
                                          (1, 0, 0, 1),
                                          (1, 0, 1, 0),
                                          (1, 0, 1, 1),
                                          (1, 1, 0, 0),
                                          (1, 1, 0, 1),
                                          (1, 1, 1, 0),
                                          (1, 1, 1, 1)]):
            checkwidgets = widget.find_all("**/CheckBox[*]")
            if index & 1:
                await checkwidgets[3].click()
            if index & 2:
                await checkwidgets[2].click()
            if index & 4:
                await checkwidgets[1].click()
            if index & 8:
                await checkwidgets[0].click()

            self.assertEqual(attr.Get(), expected)
            await ui_test.find("Property//Frame/**/.identifier=='control_state_inputs:MyBool4'").click()
            self.assertEqual(attr.Get(), (0, 0, 0, 0))
            omni.kit.undo.undo()
            self.assertEqual(attr.Get(), expected)
            omni.kit.undo.redo()
            self.assertEqual(attr.Get(), (0, 0, 0, 0))
