## 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 unittest
import carb
import omni.kit.test
from pxr import Usd, UsdGeom, Sdf, UsdShade


class TestMaterialSwap(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        from omni.kit.window.material_swap.scripts.material_swap import TEST_DATA_PATH
        self._usd_path = TEST_DATA_PATH.absolute()

    async def tearDown(self):
        pass

    async def test_material_swap(self):
        def get_bound_material(path: str) -> str:
            prim = omni.usd.get_context().get_stage().GetPrimAtPath(path)
            mat, rel = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
            if mat:
                return mat.GetPath().pathString
            return ""

        test_file_path = self._usd_path.joinpath("material_swap_test.usda").absolute()
        await omni.usd.get_context().open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        material_swap = omni.kit.window.material_swap.get_instance()
        selection = omni.usd.get_context().get_selection()

        # verify initial state
        self.assertEqual(get_bound_material('/World/studiohemisphere/Sphere'), "/World/Looks/GroundMat")
        self.assertEqual(get_bound_material('/World/studiohemisphere/Floor'), "/World/Looks/GroundMat")

        # select old material
        selection.set_selected_prim_paths(["/World/Looks/GroundMat"], True)
        material_swap._get_selectedold_fn(None)
        await omni.kit.app.get_app().next_update_async()

        # select new material
        selection.set_selected_prim_paths(["/World/Looks/PreviewSurfaceTexture"], True)
        material_swap._get_selectednew_fn(None)
        await omni.kit.app.get_app().next_update_async()

        # swap material
        material_swap._on_swap_fn(None)
        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()

        # verify change
        self.assertEqual(get_bound_material('/World/studiohemisphere/Sphere'), "/World/Looks/PreviewSurfaceTexture")
        self.assertEqual(get_bound_material('/World/studiohemisphere/Floor'), "/World/Looks/PreviewSurfaceTexture")
