# Copyright (c) 2020, 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 omni.usd
import omni.kit.app

from pxr import Usd, UsdShade, Sdf
from omni.kit.test_suite.helpers import get_test_data_path
from omni.kit.usd.layers import LayerUtils


class TestMaterialWatcher(omni.kit.test.AsyncTestCase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._test_scene = get_test_data_path(__name__, "ColorsSaltBox/SaltBox.usda")

    async def setUp(self):
        await omni.usd.get_context().new_stage_async()
        self.stage = omni.usd.get_context().get_stage()
        self.stage.GetRootLayer().subLayerPaths.append(self._test_scene)
        self.context = omni.usd.get_context()
        self.root_prim_path = Sdf.Path("/SaltBox")
        self.shader_prim_path = Sdf.Path("/SaltBox/Looks/SaltBox_Paper1/Shader")
        self.app = omni.kit.app.get_app()
        self.selection = self.context.get_selection()
        self.diffuse_constant_color_path = self.shader_prim_path.AppendProperty("inputs:diffuse_color_constant")
        self.albedo_brightness_path = self.shader_prim_path.AppendProperty("inputs:albedo_brightness")
        self.session_layer = self.stage.GetSessionLayer()
        self.stage.SetEditTarget(Usd.EditTarget(self.stage.GetRootLayer()))

        renderer = "rtx"
        if renderer not in self.context.get_attached_hydra_engine_names():
            omni.usd.add_hydra_engine(renderer, self.context)

        await self._wait()
        self.context.add_to_pending_creating_mdl_paths(str(self.shader_prim_path), True, True)
        await self._wait(2)

    async def tearDown(self):
        pass

    async def _wait(self, frames=5):
        for _ in range(frames):
            await self.app.next_update_async()

    async def test_referenced_material(self):
        await self.context.new_stage_async()
        self.stage = self.context.get_stage()
        self.session_layer = self.stage.GetSessionLayer()
        layer = Sdf.Layer.CreateAnonymous()
        self.stage.GetRootLayer().subLayerPaths.append(layer.identifier)

        # Writes materials into the sublayer to keep edit target empty.
        with Usd.EditContext(self.stage, layer):
            root_prim = self.stage.DefinePrim("/SaltBox")
            root_prim.GetReferences().AddReference(self._test_scene)

        await self._wait()

        # OM-84443: Populating material params for referenced prim will not write all back to root layer
        self.selection.set_selected_prim_paths([str(self.shader_prim_path)], False)
        await self._wait()

        shader_prim = UsdShade.Shader.Get(self.stage, self.shader_prim_path)
        self.assertTrue(root_prim)
        self.assertTrue(shader_prim)
        await self.__check_material(root_prim, shader_prim)

    async def test_material_watcher(self):
        root_prim = self.stage.GetPrimAtPath(self.root_prim_path)
        shader_prim = UsdShade.Shader.Get(self.stage, self.shader_prim_path)
        self.assertTrue(root_prim)
        self.assertTrue(shader_prim)

        # Populating material params
        self.selection.set_selected_prim_paths([str(self.shader_prim_path)], False)
        await self._wait()

        await self.__check_material(root_prim, shader_prim)

    async def __check_material(self, root_prim, shader_prim):
        # Checking params in the session layer
        shader_prim_spec = self.session_layer.GetPrimAtPath(self.shader_prim_path)
        self.assertTrue(shader_prim_spec)

        # Checkes to make sure params are not written to the current edit target after selection.
        root_layer = self.stage.GetRootLayer()
        self.assertFalse(root_layer.GetPrimAtPath(self.shader_prim_path))

        # It's overrided already in USD, so it's empty in session layer.
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_spec)
        self.assertEqual(constant_color_spec.default, None)

        constant_color_composed = self.stage.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_composed)
        self.assertEqual(constant_color_composed.Get(), (0, 0, 1))

        albedo_brightness_spec = self.session_layer.GetPropertyAtPath(self.albedo_brightness_path)
        self.assertFalse(albedo_brightness_spec.default)

        variant_set = root_prim.GetVariantSet("materials")
        self.assertTrue(variant_set)
        variant_set.SetVariantSelection("renderLow")
        await self._wait()

        # It's not overrided in USD, then the default value will be populated into session layer.
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_spec)
        self.assertEqual(constant_color_spec.default, (0.2, 0.2, 0.2))
        self.assertFalse(albedo_brightness_spec.default)

        # Sets value to root layer will remove the overrides in session layer
        constant_color_composed = self.stage.GetPropertyAtPath(self.diffuse_constant_color_path)
        constant_color_composed.Set((0.5, 0.5, 0.5))
        await self._wait()
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(not constant_color_spec or not constant_color_spec.default)

        # Removes the overrides will cause session layer to populate the default value.
        LayerUtils.remove_prim_spec(self.stage.GetRootLayer(), self.shader_prim_path)
        await self._wait()
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertEqual(constant_color_spec.default, (0.2, 0.2, 0.2))

        # Switches it back will remove the default value in session layer.
        variant_set.SetVariantSelection("renderHigh")
        await self._wait()

        # It's not overrided in USD, then the default value will be populated into session layer.
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_spec)
        self.assertEqual(constant_color_spec.default, None)
        self.assertFalse(albedo_brightness_spec.default)

        # OM-110849: Sets default value in root layer and then edits it with EditContext in session layer will not
        # been removed.
        constant_color_composed = self.stage.GetPropertyAtPath(self.diffuse_constant_color_path)
        constant_color_composed.Set((0.5, 0.5, 0.5))
        await self._wait()

        with Usd.EditContext(self.stage, self.session_layer):
            constant_color_composed.Set((0.2, 0.2, 0.2))

        await self._wait()
        constant_color_spec = self.session_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_spec)
        self.assertEqual(constant_color_spec.default, (0.2, 0.2, 0.2))

        constant_color_spec = root_layer.GetPropertyAtPath(self.diffuse_constant_color_path)
        self.assertTrue(constant_color_spec)
        self.assertEqual(constant_color_spec.default, (0.5, 0.5, 0.5))
