## 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.test
import os
import omni.kit.app
import omni.usd
from omni.kit.test.async_unittest import AsyncTestCase
from omni.kit import ui_test
from pxr import Sdf, UsdShade
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, select_prims, wait_stage_loading, delete_prim_path_children, arrange_windows


class StageMenuCreateMaterials(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows(topleft_window="Stage", topleft_height=512, topleft_width=768)
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))

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

    DATA = ["/World/Cube", "/World/Cone", "/World/Sphere", "/World/Cylinder"]

    async def test_l1_stage_menu_create_materials(self):
        stage_window = ui_test.find("Stage")
        await stage_window.focus()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        to_select = self.DATA

        # don't write to list returned by get_mdl_list as its ptr to class list
        mdl_list = await omni.kit.material.library.get_mdl_list_async(use_hidden=True)
        items_to_test = mdl_list.copy()
        items_to_test.append(("USD_Preview_Surface", "", None))
        items_to_test.append(("USD_Preview_Surface_Texture", "", None))

        for mtl_name, _, submenu in items_to_test:
            # delete any materials in looks
            await delete_prim_path_children("/World/Looks")

            # select prims
            await select_prims(to_select)

            # right click on Cube
            stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
            await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

            # click on context menu item
            menu_name = f"Create/Material/{submenu}/{mtl_name.replace('_', ' ')}" if submenu else f"Create/Material/{mtl_name.replace('_', ' ')}"
            await ui_test.select_context_menu(menu_name)

            # wait for material to load & UI to refresh
            await wait_stage_loading()

            # verify item(s)
            for prim_path in to_select:
                prim = stage.GetPrimAtPath(prim_path)
                bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                self.assertTrue(bound_material.GetPrim().IsValid() == True)
                if mtl_name in ["USD_Preview_Surface", "USD_Preview_Surface_Texture"]:
                    cmp_name = mtl_name.replace('USD_', '').replace('_', '')
                    self.assertTrue(bound_material.GetPrim().GetPrimPath().pathString == f"/World/Looks/{cmp_name}")
                else:
                    self.assertTrue(bound_material.GetPrim().GetPrimPath().pathString == f"/World/Looks/{mtl_name}")
                    prim = stage.GetPrimAtPath(f"/World/Looks/{mtl_name}/Shader")
                    asset_path = prim.GetAttribute("info:mdl:sourceAsset").Get()
                    self.assertFalse(os.path.isabs(asset_path.path))
                    self.assertTrue(os.path.isabs(asset_path.resolvedPath))
