## Copyright (c) 2018-2019, 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
from pathlib import Path
from omni.kit.widget.stage import StageWidget
from omni.kit.widget.stage.stage_icons import StageIcons
from omni.ui.tests.test_base import OmniUiTest
from pxr import Usd, UsdGeom
from omni.kit import ui_test
from omni.kit.test_suite.helpers import arrange_windows, wait_stage_loading
import omni.kit.app

CURRENT_PATH = Path(__file__).parent
REPO_PATH = CURRENT_PATH
for i in range(10):
    REPO_PATH = REPO_PATH.parent
BOWL_STAGE = REPO_PATH.joinpath("data/usd/tests/bowl.usd")


STYLE = {"Field": {"background_color": 0xFF24211F, "border_radius": 2}}


class TestStage(OmniUiTest):
    # Before running each test
    async def setUp(self):
        await arrange_windows(hide_viewport=True)

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

    async def test_general(self):
        """Testing general look of StageWidget"""
        StageIcons()._icons = {}

        window = await self.create_test_window()

        stage = Usd.Stage.Open(f"{BOWL_STAGE}")
        if not stage:
            self.fail(f"Stage {BOWL_STAGE} doesn't exist")
            return

        with window.frame:
            stage_widget = StageWidget(stage, style=STYLE)

        # 5 frames to rasterize the icons
        for _ in range(5):
            await omni.kit.app.get_app().next_update_async()

        await self.finalize_test()

    async def test_visibility(self):
        """Testing visibility of StageWidget"""
        StageIcons()._icons = {}

        window = await self.create_test_window()

        stage = Usd.Stage.CreateInMemory("test.usd")

        with window.frame:
            stage_widget = StageWidget(stage, style=STYLE)

        UsdGeom.Mesh.Define(stage, "/A")
        UsdGeom.Mesh.Define(stage, "/A/B")
        UsdGeom.Mesh.Define(stage, "/C").MakeInvisible()

        # 5 frames to rasterize the icons
        for _ in range(5):
            await omni.kit.app.get_app().next_update_async()

        await self.finalize_test()

    async def test_dynamic(self):
        """Testing the ability to watch the stage"""
        StageIcons()._icons = {}

        window = await self.create_test_window()

        stage = Usd.Stage.CreateInMemory("test.usd")

        with window.frame:
            stage_widget = StageWidget(stage, style=STYLE)

        UsdGeom.Mesh.Define(stage, "/A")

        # Wait one frame to be sure other objects created after initialization
        await omni.kit.app.get_app().next_update_async()

        UsdGeom.Mesh.Define(stage, "/B")
        mesh = UsdGeom.Mesh.Define(stage, "/C")

        # Wait one frame to be sure other objects created after initialization
        await omni.kit.app.get_app().next_update_async()

        mesh.MakeInvisible()
        UsdGeom.Mesh.Define(stage, "/D")

        # 5 frames to rasterize the icons
        for _ in range(5):
            await omni.kit.app.get_app().next_update_async()

        await self.finalize_test()

    async def test_instanceable(self):
        """Testing the ability to watch the instanceable state"""
        StageIcons()._icons = {}

        window = await self.create_test_window()

        stage = Usd.Stage.CreateInMemory("test.usd")

        with window.frame:
            stage_widget = StageWidget(stage, style=STYLE)

        UsdGeom.Xform.Define(stage, "/A")
        UsdGeom.Mesh.Define(stage, "/A/Shape")

        prim = UsdGeom.Xform.Define(stage, "/B").GetPrim()
        prim.GetReferences().AddInternalReference("/A")

        # Wait one frame to be sure other objects created after initialization
        await omni.kit.app.get_app().next_update_async()

        stage_widget.expand("/B")

        await omni.kit.app.get_app().next_update_async()

        prim.SetInstanceable(True)

        # 5 frames to rasterize the icons
        for _ in range(5):
            await omni.kit.app.get_app().next_update_async()

        await self.finalize_test()
