## 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.
##
from sys import path
import omni.kit.app
import omni.kit.commands
import omni.kit.undo
import omni.kit.test
import omni.timeline
import carb.settings
from .test_hydra_common import RtxHydraTest, RtxHydraMaterialsTest, RtxHydraInstancerTest
from .test_common import testSettings, postLoadTestSettings, wait_for_update, USD_DIR
from pxr import Gf, Sdf, UsdGeom, UsdRender, UsdShade
import shutil, os, pathlib


class RtxHydraMaterialTest(RtxHydraTest):
    TEST_PATH = "hydra/materials"
    PRIM_PATH = "/World/box"
    DATA_PATH = "hydra/materials/reload"
    TEST_CLASS = ""

    # need to use the real path in order to have the filewatching work properly
    # -> client library ticket: https://nvidia-omniverse.atlassian.net/browse/CC-168
    def getDataPath(self, relativePath) -> str:
        data_path : str = USD_DIR.joinpath(self.DATA_PATH, relativePath)
        return pathlib.Path(os.path.normpath(data_path)).resolve()

    # helper to override the test files
    # in case we want to do more or use another function
    def copyTestFile(self, source : str, target : str):
        shutil.copy(self.getDataPath(source), self.getDataPath(target))

    # delete the temporary file created
    def cleanupTestFile(self, target : str):
        os.remove(self.getDataPath(target))

    # basic structure of a test case init
    async def initTest(self, caseName: str):
        goldenImage = f"reload-{self.TEST_CLASS}-{caseName}.png"

        # compare
        await wait_for_update(wait_frames=10)
        await self.capture_and_compare(self.TEST_PATH, goldenImage)

    # basic structure of a test case
    async def changeMdlAndTest(self, caseName: str, moduleToChange: str, threshold=None):
        sourceFile = f"{self.TEST_CLASS}/{caseName}.mdl"
        targetFile = f"{self.TEST_CLASS}/{moduleToChange}.mdl"
        goldenImage = f"reload-{self.TEST_CLASS}-{caseName}.png"

        if threshold is None:
            threshold = self.THRESHOLD

        # change MDL
        self.copyTestFile(sourceFile, targetFile)
        # compare
        await wait_for_update(wait_frames=10)
        await self.capture_and_compare(self.TEST_PATH, goldenImage, threshold)


class TestRtxHydraMaterialsReloadBasic(RtxHydraMaterialTest):
    MEDIUM_THRESHOLD = 3.0e-5

    async def setUp(self):
        self.TEST_CLASS = "basic"
        await self.setUp_internal()
        self.set_settings(testSettings)

        # prepare initial MDL module
        self.copyTestFile("basic/A_init.mdl", "basic/A.mdl")

        # load the test scene
        scenePath = self.getDataPath("basic/XYZ.usda")
        print(f"Scene Path: {scenePath}")
        super().open_usd(scenePath)
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    async def test_touch_no_change(self):
        """
        Test hydra materials - reload basic - touch, no change
        """
        # make sure we start with the right material
        await self.initTest("A_init")

        # touch the file by copying the same file again
        await self.changeMdlAndTest("A_init", "A", self.MEDIUM_THRESHOLD)

        # remove the temp file
        self.cleanupTestFile(f"{self.TEST_CLASS}/A.mdl")

    async def test_change_body(self):
        """
        Test hydra materials - reload basic - change body
        """
        # make sure we start with the right material
        await self.initTest("A_init")

        # make multiple changes in a row
        await self.changeMdlAndTest("A_change_body_1", "A", self.MEDIUM_THRESHOLD)
        await self.changeMdlAndTest("A_change_body_2", "A", self.MEDIUM_THRESHOLD)

        # remove the temp file
        self.cleanupTestFile(f"{self.TEST_CLASS}/A.mdl")

class TestRtxHydraMaterialsReloadSignature(RtxHydraMaterialTest):
    async def setUp(self):
        self.TEST_CLASS = "signature"
        await self.setUp_internal()
        self.set_settings(testSettings)

        # prepare initial MDL module
        self.copyTestFile("signature/A_init.mdl", "signature/A.mdl")

        # load the test scene
        scenePath = self.getDataPath("signature/A.usda")
        print(f"Scene Path: {scenePath}")
        super().open_usd(scenePath)
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)


    async def test_change_signature(self):
        """
        Test hydra materials - reload signature - change signature
        """
        # make sure we start with the right material
        await self.initTest("A_init")

        # run test cases
        await self.changeMdlAndTest("A_add_parameter", "A")
        await self.changeMdlAndTest("A_rename_parameter", "A")
        await self.changeMdlAndTest("A_move_parameter", "A")
        await self.changeMdlAndTest("A_change_parameter_default", "A")
        await self.changeMdlAndTest("A_change_parameter_type", "A")
        await self.changeMdlAndTest("A_remove_parameter", "A")

        # remove the temp file
        self.cleanupTestFile(f"{self.TEST_CLASS}/A.mdl")

class TestRtxHydraMaterialsReloadGraph(RtxHydraMaterialTest):
    async def setUp(self):
        self.TEST_CLASS = "graph"
        await self.setUp_internal()
        self.set_settings(testSettings)

        # prepare initial MDL module
        self.copyTestFile("graph/A_init.mdl", "graph/A.mdl")

        # load the test scene
        scenePath = self.getDataPath("graph/GraphA.usda")
        print(f"Scene Path: {scenePath}")
        super().open_usd(scenePath)
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    # Unstable via OM-53607
    async def test_UNSTABLE_change_function_body_UNSTABLE(self):
        """
        Test hydra materials - reload graph - change function body
        """
        # make sure we start with the right material
        await self.initTest("A_init")

        # run test cases
        await self.changeMdlAndTest("A_change_function_body", "A")
        await self.changeMdlAndTest("A_change_function_param_defaults", "A")
        await self.changeMdlAndTest("A_add_function_param", "A")
        await self.changeMdlAndTest("A_move_function_param", "A")
        await self.changeMdlAndTest("A_remove_function_param", "A")

        # remove the temp file
        self.cleanupTestFile(f"{self.TEST_CLASS}/A.mdl")

class TestRtxHydraMaterialsDebugMode(RtxHydraMaterialTest):
    async def setUp(self):
        self.TEST_CLASS = "debugMode"
        await self.setUp_internal()
        self.set_settings(testSettings)

    async def test_whiteMode(self):
        """
        Test hydra materials - toggling white mode
        """
        super().open_usd("hydra/materials/two_spheres.usda")
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "materials_whiteMode_off.png", 1e-4)

        self.set_settings({"/rtx/debugMaterialType": 0})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "materials_whiteMode_on.png", 1e-4)

        shader = self.ctx.get_stage().GetPrimAtPath("/World/Looks/PreviewSurface/Shader")
        attr = shader.GetAttribute("inputs:excludeFromWhiteMode")
        attr.Set(1)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "materials_whiteMode_on_withExclude.png", 1e-4)

        attr.Set(0)
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "materials_whiteMode_on_withoutExclude.png", 1e-4)

        self.set_settings({"/rtx/debugMaterialType": -1})
        await wait_for_update()
        await self.capture_and_compare(self.TEST_PATH, "materials_whiteMode_off_again.png", 1e-4)

class TestRtxHydraMaterialStdModules(RtxHydraTest):
    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        super().open_usd("hydra/materials/stdmodules/test_std_modules.usda")
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    async def test_load_stdmodules(self):
        await wait_for_update(wait_frames=10)
        await self.capture_and_compare("hydra/materials", "stdmodules.png")
