## 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.test
import math
import omni.usd
import omni.kit.app
from pxr import UsdGeom
from omni.kit.test.async_unittest import AsyncTestCase
from omni.kit import ui_test
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, wait_stage_loading, get_prims, arrange_windows


class VisibilityToggleUsdStage(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows("Stage", 512)
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))

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

    async def test_l1_eye_visibility_icon(self):
        await ui_test.find("Content").focus()
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        await wait_stage_loading()

        def verify_prim_state(expected):
            stage = omni.usd.get_context().get_stage()
            for prim_path in [prim.GetPath().pathString for prim in stage.TraverseAll() if not omni.usd.is_hidden_type(prim)]:
                if not "Looks" in prim_path:
                    self.assertEqual(UsdGeom.Imageable(stage.GetPrimAtPath(prim_path)).ComputeVisibility(), expected[prim_path])

        # veirfy default state
        verify_prim_state({"/World": "inherited", "/World/defaultLight": "inherited", "/World/Cone": "inherited", "/World/Cube": "inherited", "/World/Sphere": "inherited", "/World/Cylinder": "inherited"})

        # build table of eye buttons
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        widgets = {}
        for w in stage_widget.find_all("**/Label[*]"):
            widget_name = w.widget.text if w.widget.text != "World (defaultPrim)" else "World"
            for wtb in stage_widget.find_all(f"**/ToolButton[*]"):
                if math.isclose(wtb.widget.screen_position_y, w.widget.screen_position_y):
                    widgets[widget_name] = wtb
                    break

        # click world eye & verify
        await widgets["World"].click()
        verify_prim_state({"/World": "invisible", "/World/defaultLight": "invisible", "/World/Cone": "invisible", "/World/Cube": "invisible", "/World/Sphere": "invisible", "/World/Cylinder": "invisible"})
        await widgets["World"].click()
        verify_prim_state({"/World": "inherited", "/World/defaultLight": "inherited", "/World/Cone": "inherited", "/World/Cube": "inherited", "/World/Sphere": "inherited", "/World/Cylinder": "inherited"})

        # click individual prima eye & verify
        for prim_name in ["defaultLight", "Cone", "Cube", "Sphere", "Cylinder"]:
            expected = {"/World": "inherited", "/World/defaultLight": "inherited", "/World/Cone": "inherited", "/World/Cube": "inherited", "/World/Sphere": "inherited", "/World/Cylinder": "inherited"}
            verify_prim_state(expected)
            await widgets[prim_name].click()
            expected[f"/World/{prim_name}"] = "invisible"
            verify_prim_state(expected)
            await widgets[prim_name].click()
