# NOTE:
#   omni.kit.test - std python's unittest module with additional wrapping to add suport for async/await tests
#   For most things refer to unittest docs: https://docs.python.org/3/library/unittest.html
import omni.kit.test

# Import extension python module we are testing with absolute import path, as if we are external user (other extension)
import omni.mdl.usd_converter

import omni.usd

import shutil
import pathlib
import os
import carb
import carb.settings

# Having a test class derived from omni.kit.test.AsyncTestCase declared on the root of module
# will make it auto-discoverable by omni.kit.test
class Test(omni.kit.test.AsyncTestCaseFailOnLogError):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._test_data = str(pathlib.Path(__file__).parent.joinpath("data"))
        self._test_data_usd = str(pathlib.Path(self._test_data).joinpath("usd"))
        self._test_data_mdl = str(pathlib.Path(self._test_data).joinpath("mdl"))

    # Before running each test
    async def setUp(self):
        pass

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

    # Actual test, notice it is "async" function, so "await" can be used if needed
    async def test_mdl_to_usd(self):

        # Disable standard path to ensure reproducible results locally and in TC
        NO_STD_PATH = "/app/mdl/nostdpath"
        settings = carb.settings.get_settings()
        settings.set_bool(NO_STD_PATH, True)

        filter_import = False
        tokens = omni.mdl.usd_converter.find_tokens('mdl::ZA0OmniGlass_2Emdl::OmniGlass::converted_27', filter_import)
        result = (tokens == {'ZA0OmniGlass_2Emdl', 'converted_27', 'mdl', 'OmniGlass'})
        self.assertEqual(result, True)

        filter_import = True
        tokens = omni.mdl.usd_converter.find_tokens('import ::state::normal;\n', filter_import)
        result = (tokens == {'state', 'normal'})
        self.assertEqual(result, True)

        filter_import = False
        tokens = omni.mdl.usd_converter.find_tokens('    anno::author("NVIDIA CORPORATION"),\n', filter_import)
        result = (tokens == {'anno', 'author'})
        self.assertEqual(result, True)

        filter_import = True
        tokens = omni.mdl.usd_converter.find_tokens('    anno::author("NVIDIA CORPORATION"),\n', filter_import)
        result = (len(tokens) == 0)
        self.assertEqual(result, True)

        with omni.mdl.usd_converter.TemporaryDirectory() as temp_dir:
            test_search_path = self._test_data_mdl
            fn = "core_definitions.usda"

            ADD_MDL_PATH = "/app/mdl/additionalUserPaths"
            settings = carb.settings.get_settings()
            mdl_custom_paths: List[str] = settings.get(ADD_MDL_PATH) or []
            mdl_custom_paths.append(test_search_path)
            mdl_custom_paths.append(temp_dir)
            mdl_custom_paths = list(set(mdl_custom_paths))
            settings.set_string_array(ADD_MDL_PATH, mdl_custom_paths)

            result = omni.mdl.usd_converter.mdl_to_usd(
                moduleName="nvidia/core_definitions.mdl",
                targetFolder=temp_dir,
                targetFilename=fn,
                output=omni.mdl.usd_converter.mdl_usd.OutputType.MATERIAL_AND_GEOMETRY)
            self.assertEqual(result, True)

            # fn = "tutorials.usda"
            # # WORKAROUND: Copy test file to temp folder, NeurayLib does not handle folder containing '*.mdl.*' in their name
            # source_file = str(pathlib.Path(self._test_data_mdl).joinpath('nvidia/sdk_examples/tutorials.mdl'))
            # shutil.copy2(source_file, temp_dir)
            # module = str(pathlib.Path(temp_dir).joinpath("tutorials.mdl"))
            # if os.path.exists(module):
            #     result = omni.mdl.usd_converter.mdl_to_usd(
            #         moduleName=module,
            #         targetFolder=temp_dir,
            #         targetFilename=fn,
            #         searchPath=test_search_path)
            #     self.assertEqual(result, True)

            # fn = "test_types.usda"
            # WORKAROUND: Copy test file to temp folder, NeurayLib does not handle folder containing '*.mdl.*' in their name
            # source_file = str(pathlib.Path(self._test_data_mdl).joinpath('test/test_types.mdl'))
            # shutil.copy2(source_file, temp_dir)
            # module = str(pathlib.Path(temp_dir).joinpath("test_types.mdl"))
            # if os.path.exists(module):
            #     result = omni.mdl.usd_converter.mdl_to_usd(
            #         moduleName=module,
            #         targetFolder=temp_dir,
            #         targetFilename=fn,
            #         searchPath=test_search_path)
            #     self.assertEqual(result, True)

            tutorial_scene = str(pathlib.Path(self._test_data_usd).joinpath("tutorials.usda"))
            if os.path.exists(tutorial_scene):
                result = await omni.mdl.usd_converter.usd_prim_to_mdl(
                    tutorial_scene,
                    "/example_material",
                    test_search_path,
                    forceNotOV=True)
                self.assertEqual(result, True)

                result = await omni.mdl.usd_converter.usd_prim_to_mdl(
                    tutorial_scene,
                    "/example_modulemdl_material_examples",
                    test_search_path,
                    forceNotOV=True)
                self.assertEqual(result, True)

                result = await omni.mdl.usd_converter.usd_prim_to_mdl(
                    tutorial_scene,
                    "/wave_gradient",
                    test_search_path,
                    forceNotOV=True)
                self.assertEqual(result, True)

            omni.mdl.usd_converter.test_mdl_prim_to_usd('/root')

            simple_scene = str(pathlib.Path(self._test_data_usd).joinpath("scene.usda"))
            if os.path.exists(simple_scene):
                test_prim = '/World/Looks/Material'
                
                stage = omni.mdl.usd_converter.mdl_usd.Usd.Stage.Open(simple_scene)
                prim = stage.GetPrimAtPath(test_prim)
                omni.mdl.usd_converter.mdl_prim_to_usd(stage, prim)

                omni.mdl.usd_converter.mdl_usd.get_examples_search_path()

                mdl_tmpfile = str(pathlib.Path(temp_dir).joinpath("tmpout.mdl"))
                result = await omni.mdl.usd_converter.export_to_mdl(mdl_tmpfile, prim, forceNotOV=True)
                self.assertEqual(result, True)

            # Example: 'mdl::OmniSurface::OmniSurfaceBase::OmniSurfaceBase'
            # Return: 'OmniSurface/OmniSurfaceBase.mdl'
            omni.mdl.usd_converter.mdl_usd.prototype_to_source_asset('mdl::OmniSurface::OmniSurfaceBase::OmniSurfaceBase')

            omni.mdl.usd_converter.mdl_usd.parse_ov_resource_name('mdl::resolvedPath_texture_raw')
