#!/usr/bin/env python3

import omni.kit.commands
import omni.kit.test
import omni.usd
from omni.rtx.tests import RtxTest, testSettings, postLoadTestSettings
from omni.rtx.tests.test_common import wait_for_update
from omni.kit.test_helpers_gfx.compare_utils import ComparisonMetric
import carb

from pathlib import Path
import os

EXTENSION_DIR = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))
TESTS_DIR = EXTENSION_DIR.joinpath('data', 'tests')
USD_DIR = TESTS_DIR.joinpath('usd')
GOLDEN_IMAGES_DIR = TESTS_DIR.joinpath('golden')
OUTPUTS_DIR = Path(omni.kit.test.get_test_output_path())

# This class is auto-discoverable by omni.kit.test
class MtlxRenderTest(RtxTest):

    WINDOW_SIZE = (1280, 768)
    THRESHOLD = 5e-5

    # common settings used for all renderers
    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)

    # Overridden with custom paths
    async def capture_and_compare(self,
                                  renderer: str,
                                  img_subdir: Path = "",
                                  test_name=None,
                                  threshold=THRESHOLD,
                                  metric: ComparisonMetric = ComparisonMetric.MEAN_ERROR_SQUARED):
        golden_img_dir = GOLDEN_IMAGES_DIR.joinpath(img_subdir)
        output_img_dir = OUTPUTS_DIR.joinpath(img_subdir)
        golden_img_name = f"{test_name}_{renderer}.png"

        return await self._capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir, metric)

    # Open the scene, overridden with custom paths
    async def open_usd(self, usdSubpath: Path):
        path = USD_DIR.joinpath(usdSubpath)
        # This settings should be set before stage opening/creation
        self.set_settings(testSettings)
        # Actually open the stage
        self.ctx.open_stage(str(path))
        await omni.kit.app.get_app().next_update_async()
        # Settings that should override settings that was set on stage opening/creation
        self.set_settings(postLoadTestSettings)

    # Close the scene
    def close_usd(self):
        self.ctx.close_stage()

    # do string replacements and save a new file
    def prepare_temp_file(self, srcFilename, destFilename, replacementMap):
        with open(srcFilename, 'r') as srcFile, \
             open(destFilename, 'w') as destFile:

            src = srcFile.read()
            for key, value in replacementMap.items():
                src = src.replace(key, value)
            destFile.write(src)

    # base test script
    async def run_image_test(self, usd_file: str, test_name: str = None):
        # load the scene
        if usd_file:
            await self.open_usd(usd_file)

        # wait for rendering and capture
        await wait_for_update(wait_frames=128+10)
        
        # golden image path need the renderer
        settings = carb.settings.get_settings()
        renderer:str = settings.get("/renderer/active")
        await self.capture_and_compare(renderer = renderer, test_name = test_name)

        # close the scene to also test clean-up
        # self.close_usd() # getting `[Error] [omni.usd] Stage opening or closing already in progress!!`
        await wait_for_update()

    # ==============================================================================================
    # The tests
    # ==============================================================================================

    # StandardSurface presets that ship with the MaterialX SDK
    # ----------------------------------------------------------------------------------------------
    async def test_standardsurface_composition(self):
        await self.run_image_test('StandardSurface/Composition.usda', 'mtlx_standardsurface_composition')

    # A selection of the materials published by AMD on https://matlib.gpuopen.com
    # ----------------------------------------------------------------------------------------------
    async def test_amd_composition(self):
        await self.run_image_test('AMD/Composition.usda', 'mtlx_amd_composition')

    # Open Chess Set as released by the USD Work Group
    # ----------------------------------------------------------------------------------------------
    async def test_open_chess_set(self):
        await self.run_image_test('OpenChessSet/chess_set_light_camera.usda', 'mtlx_open_chess_set')

    # MaterialX tests published in the USD workgroup on https://github.com/usd-wg/assets
    # Added a single root prim to each test scene and recomposed them into one test case.
    # ----------------------------------------------------------------------------------------------
    async def test_usd_wg_assets_composition(self):
        try:
            # replace the token by an absolute file path in order to test absolute paths
            sceneDir = USD_DIR.joinpath('usd-wg-assets')
            tempSceneFile = sceneDir.joinpath('basicTextured_flatten.usda')
            self.prepare_temp_file(
                srcFilename = sceneDir.joinpath('basicTextured_flatten_template.usda'),
                destFilename = tempSceneFile,
                replacementMap = {"${USD_DIR}":f"{sceneDir}"})
            tempMtlxFile = sceneDir.joinpath('standard_surface_brass_tiled_absolute_paths.mtlx')
            self.prepare_temp_file(
                srcFilename = sceneDir.joinpath('standard_surface_brass_tiled_absolute_paths_template.mtlx'),
                destFilename = tempMtlxFile,
                replacementMap = {"${USD_DIR}":f"{sceneDir}"})
            # run the test
            await self.run_image_test('usd-wg-assets/Composition.usda', 'mtlx_usd-wg-assets_composition')
        finally:
            # remove the temp files
            if tempSceneFile.exists():
                tempSceneFile.unlink()
            if tempMtlxFile.exists():
                tempMtlxFile.unlink()
