# *****************************************************************************
# Copyright 2021 NVIDIA Corporation. All rights reserved.
# *****************************************************************************

import omni.ext
import carb
import os
import shutil
import tempfile
import carb.settings

from omni.mdl import pymdlsdk
from omni.mdl import pymdl
import omni.mdl.neuraylib
from . import mdl_usd
import asyncio
from pxr import Usd
from pxr import UsdShade

MDL_AUTOGEN_PATH = "${data}/shadergraphs/mdl_usd"


class TemporaryDirectory:
    def __init__(self):
        self.path = None

    def __enter__(self):
        self.path = tempfile.mkdtemp()
        return self.path

    def __exit__(self, type, value, traceback):
        # Remove temporary data created
        shutil.rmtree(self.path)


def add_search_path_to_system_path(searchPath: str):
    if not searchPath == None and not searchPath == "":
        MDL_SYSTEM_PATH = "/app/mdl/additionalSystemPaths"
        settings = carb.settings.get_settings()
        mdl_custom_paths: List[str] = settings.get(MDL_SYSTEM_PATH) or []
        mdl_custom_paths.append(searchPath)
        mdl_custom_paths = list(set(mdl_custom_paths))
        settings.set_string_array(MDL_SYSTEM_PATH, mdl_custom_paths)


# Functions and vars are available to other extension as usual in python: `omni.mdl.usd_converter.mdl_to_usd(x)`
#
# mdl_to_usd(moduleName, targetFolder, targetFilename)
#   moduleName: module to convert (example: "nvidia/core_definitions.mdl")
#   searchPath: MDL search path to be able to load MDL modules referenced by moduleName
#   targetFolder: Destination folder for USD stage (default = "${data}/shadergraphs/mdl_usd")
#   targetFilename: Destination stage filename (default is module name, example: "core_definitions.usda")
#   output: What to output:
#       a shader: omni.mdl.usd_converter.mdl_usd.OutputType.SHADER
#       a material: omni.mdl.usd_converter.mdl_usd.OutputType.MATERIAL
#       geometry and material: omni.mdl.usd_converter.mdl_usd.OutputType.MATERIAL_AND_GEOMETRY
#   nestedShaders: Do we want nested shaders or flat (default = False)
#
def mdl_to_usd(
        moduleName: str,
        searchPath: str = None,
        targetFolder: str = MDL_AUTOGEN_PATH,
        targetFilename: str = None,
        output: mdl_usd.OutputType = mdl_usd.OutputType.SHADER,
        nestedShaders: bool = False):
    print(f"[omni.mdl.usd_converter] mdl_to_usd was called with {moduleName}")
    # acquire neuray instance from OV
    ovNeurayLib = omni.mdl.neuraylib.get_neuraylib()
    ovNeurayLibHandle = ovNeurayLib.getNeurayAPI()
    # feed the neuray instance into the python binding
    neuray: pymdlsdk.INeuray = pymdlsdk.attach_ineuray(ovNeurayLibHandle)
    neurayStatus: pymdlsdk.INeuray.Status = neuray.get_status()
    print(f"Neuray Status: {neurayStatus}")

    # Set carb.settings with new MDL search path
    add_search_path_to_system_path(searchPath)

    # Select a view on the database used for the rtx renderer for.
    # It should be fine to use this one always. Hydra will deal with applying the changes to the
    # other renderer. It would be possible to use an own space entirely but this would double module loading.
    # In the long, we want to load modules only into one scope.
    dbScopeName = "rtx_scope"
    # we need to load modules to OV using the omni.mdl.neuraylib
    # on the c++ side this is async, here it is blocking!
    print(f"createMdlModule called with: {moduleName}")
    ovModule = ovNeurayLib.createMdlModule(moduleName)
    print(f"CoreDefinitions: {ovModule.valid()}")
    if ovModule.valid():
        print(f" dbScopeName: {ovModule.dbScopeName}")
        print(f" dbName: {ovModule.dbName}")
        print(f" qualifiedName: {ovModule.qualifiedName}")
    else:
        print(f" createMdlModule failed : {moduleName}")

    # after the module is loaded we create a new transaction that can see the loaded module
    ovTransactionReadHandle = ovNeurayLib.createReadingTransaction(dbScopeName)

    trans: pymdlsdk.ITransaction = pymdlsdk.attach_itransaction(ovTransactionReadHandle)
    print(f"Transaction Open: {trans.is_open()}")

    # try out the high level binding
    module = pymdl.Module._fetchFromDb(trans, ovModule.dbName)
    if module:
        print(f"MDL Qualified Name: {module.mdlName}")
        print(f"MDL Simple Name: {module.mdlSimpleName}")
        print(f"MDL DB Name: {module.dbName}")
        print(f"MDL Module filename: {module.filename}")
    else:
        print("MDL Module: None")
        print(f"Failed to convert: {moduleName}")
        return False

    try:
        stage = mdl_usd.Usd.Stage.CreateInMemory()
    except:
        trans.abort()
        trans = None
        return False
        
    stage.SetMetadata('comment', 'MDL to USD conversion')

    # create context for the conversion of the scene
    context = mdl_usd.ConverterContext()
    context.set_neuray(neuray)
    context.set_transaction(trans)
    context.set_stage(stage)
    context.mdl_to_usd_output = output
    context.mdl_to_usd_output_material_nested_shaders = nestedShaders
    context.ov_neuray = ovNeurayLib

    mdl_usd.module_to_stage(context, module)
    # Workaround the issue creating stage under "${data}/shadergraphs/mdl_usd"
    # by creating the stage in a temp folder and copying it to dest folder
    with TemporaryDirectory() as temp_dir:
        unmangle_helper = mdl_usd.Unmangle(context)
        # Demangle instance name
        (unmangled_flag, simple_name) = unmangle_helper.unmangle_mdl_identifier(module.dbName)
        simple_name = simple_name.replace("::", "_")
        filename = os.path.join(temp_dir, simple_name + ".usda")
        if targetFilename is not None:
            if targetFilename[-4:] == ".usd" or targetFilename[-5:] == ".usda":
                filename = os.path.join(temp_dir, targetFilename)
            else:
                filename = os.path.join(temp_dir, targetFilename + ".usda")
        stage.GetRootLayer().Export(filename)

        # Copy file to destination folder
        path = targetFolder
        token = carb.tokens.get_tokens_interface()
        mdlUSDPath = token.resolve(path)
        dest = os.path.abspath(mdlUSDPath)
        # Create folder if it does not exist
        if not os.path.exists(dest):
            os.makedirs(dest)
        try:
            shutil.copy2(filename, dest)
            print(f"Stage saved as: {os.path.join(dest, os.path.basename(filename))}")
        except:
            print(f"Failed to save stage as: {os.path.join(dest, os.path.basename(filename))}")
            pass

    try:
        omni.kit.window.material_graph.GraphExtension.refresh_compounds()
    except:
        pass

    ovNeurayLib.destroyMdlModule(ovModule)

    # since we have been reading only, abort
    trans.abort()
    trans = None

    return True


def test_mdl_prim_to_usd(primPath: str):
    stage = omni.usd.get_context().get_stage()
    if stage:
        prim = stage.GetPrimAtPath(primPath)
        if prim:
            mdl_prim_to_usd(stage, prim)


def mdl_prim_to_usd(stage: Usd.Stage, prim: Usd.Prim):
    # Only handle material for the time beeing
    if not UsdShade.Material(prim):
        return

    # acquire neuray instance from OV
    ovNeurayLib = omni.mdl.neuraylib.get_neuraylib()
    ovNeurayLibHandle = ovNeurayLib.getNeurayAPI()
    # feed the neuray instance into the python binding
    neuray: pymdlsdk.INeuray = pymdlsdk.attach_ineuray(ovNeurayLibHandle)
    neurayStatus: pymdlsdk.INeuray.Status = neuray.get_status()
    print(f"Neuray Status: {neurayStatus}")

    # after the module is loaded we create a new transaction that can see the loaded module
    dbScopeName = "rtx_scope"
    ovTransactionReadHandle = ovNeurayLib.createReadingTransaction(dbScopeName)

    trans: pymdlsdk.ITransaction = pymdlsdk.attach_itransaction(ovTransactionReadHandle)
    print(f"Transaction Open: {trans.is_open()}")

    # create context for the conversion of the scene
    context = mdl_usd.ConverterContext()
    context.set_neuray(neuray)
    context.set_transaction(trans)
    context.set_stage(stage)
    context.mdl_to_usd_output = mdl_usd.OutputType = mdl_usd.OutputType.MATERIAL
    context.mdl_to_usd_output_material_nested_shaders = False
    context.ov_neuray = ovNeurayLib

    mdl_usd.mdl_prim_to_usd(context, stage, prim)   

    # since we have been reading only, abort
    trans.abort()
    trans = None

    return True


# usd_prim_to_mdl(usdStage, usdPrim, searchPath)
#   usdStage: Input stage containing the Prim to convert
#   usdPrim: Prim to convert (default = not set, use stage default Prim)
#   searchPath: MDL search path to be able to load MDL modules referenced by usdPrim (default = not set)
#   forceNotOV: If set to True do not use the OV NeurayLib for conversion. Use specific code. (default = False)
#
async def usd_prim_to_mdl(usdStage: str, usdPrim: str = None, searchPath: str = None, forceNotOV: bool = False):
    print(f"[omni.mdl.usd_converter] usd_prim_to_mdl was called with {usdStage} / {usdPrim} / {searchPath} / {forceNotOV}")
    rtncode = True
    # acquire neuray instance from OV
    ovNeurayLib = omni.mdl.neuraylib.get_neuraylib()
    ovNeurayLibHandle = ovNeurayLib.getNeurayAPI()
    # feed the neuray instance into the python binding
    neuray: pymdlsdk.INeuray = pymdlsdk.attach_ineuray(ovNeurayLibHandle)
    neurayStatus: pymdlsdk.INeuray.Status = neuray.get_status()
    print(f"Neuray Status: {neurayStatus}")

    # Set carb.settings with new MDL search path
    add_search_path_to_system_path(searchPath)

    # Select a view on the database used for the rtx renderer for.
    # It should be fine to use this one always. Hydra will deal with applying the changes to the
    # other renderer. It would be possible to use an own space entirely but this would double module loading.
    # In the long, we want to load modules only into one scope.
    dbScopeName = "rtx_scope"

    # create a new transaction
    ovTransactionReadHandle = ovNeurayLib.createReadingTransaction(dbScopeName)

    trans: pymdlsdk.ITransaction = pymdlsdk.attach_itransaction(ovTransactionReadHandle)
    print(f"Transaction Open: {trans.is_open()}")

    try:
        stage = mdl_usd.Usd.Stage.Open(usdStage)
    except:
        rtncode = False
    
    if rtncode:
        if usdPrim == None:
            try:
                # Determine default Prim
                usdPrim = stage.GetDefaultPrim().GetPath()
            except:
                pass

        if usdPrim == None:
            print(f"[omni.mdl.usd_converter] error: No prim specified and no default prim in stage")
            rtncode = False

        if rtncode:
            # create context for the conversion of the scene
            context = mdl_usd.ConverterContext()
            context.set_neuray(neuray)
            context.set_transaction(trans)
            context.set_stage(stage)
            if not forceNotOV:
                context.ov_neuray = ovNeurayLib

            inst_name = await mdl_usd.convert_usd_to_mdl(context, usdPrim, None)

            rtncode = (inst_name is not None)

            # transaction might have changed internally
            trans = context.transaction

    if not rtncode:
        print(f"[omni.mdl.usd_converter] error: Conversion failure")

    trans.abort()

    trans = None

    return rtncode


# test_export_to_mdl(usdStage, usdPrim, searchPath)
#   usdStage: Input stage containing the Prim to convert
#   usdPrim: Prim to convert (default = not set, use stage default Prim)
#   searchPath: MDL search path to be able to load MDL modules referenced by usdPrim (default = not set)
#
async def test_export_to_mdl(usdStage: str, usdPrim: str = None, searchPath: str = None):  # pragma: no cover

    # Set carb.settings with new MDL search path
    add_search_path_to_system_path(searchPath)

    stage = None
    try:
        stage = mdl_usd.Usd.Stage.Open(usdStage)
    except:
        return False
    
    if usdPrim == None:
        try:
            # Determine default Prim
            usdPrim = stage.GetDefaultPrim().GetPath()
        except:
            pass

    if usdPrim == None:
        print(f"[omni.mdl.usd_converter] error: No prim specified and no default prim in stage")
        return False

    prim = stage.GetPrimAtPath(usdPrim)

    path = "c:/temp/new_module.mdl"

    return asyncio.ensure_future(export_to_mdl(path, prim))


# export_to_mdl(path, prim)
#   path: Output file name
#   prim: Prim to convert
#
#   Note: The MDL modules referenced by the prim must be in the MDL searchPath
#
async def export_to_mdl(path: str, prim: mdl_usd.Usd.Prim, forceNotOV: bool = False):

    if prim == None:
        print(f"[omni.mdl.usd_converter] error: No prim specified and no default prim in stage")
        return False
    
    # acquire neuray instance from OV
    ovNeurayLib = omni.mdl.neuraylib.get_neuraylib()
    ovNeurayLibHandle = ovNeurayLib.getNeurayAPI()
    # feed the neuray instance into the python binding
    neuray: pymdlsdk.INeuray = pymdlsdk.attach_ineuray(ovNeurayLibHandle)
    neurayStatus: pymdlsdk.INeuray.Status = neuray.get_status()
    print(f"Neuray Status: {neurayStatus}")

    # Select a view on the database used for the rtx renderer for.
    # It should be fine to use this one always. Hydra will deal with applying the changes to the
    # other renderer. It would be possible to use an own space entirely but this would double module loading.
    # In the long, we want to load modules only into one scope.
    dbScopeName = "rtx_scope"
  
    # create a new transaction
    ovTransactionReadHandle = ovNeurayLib.createReadingTransaction(dbScopeName)

    trans: pymdlsdk.ITransaction = pymdlsdk.attach_itransaction(ovTransactionReadHandle)
    print(f"Transaction Open: {trans.is_open()}")

    stage = prim.GetStage()
    
    # create context for the conversion of the scene
    context = mdl_usd.ConverterContext()
    context.set_neuray(neuray)
    context.set_transaction(trans)
    context.set_stage(stage)
    if not forceNotOV:
        context.ov_neuray = ovNeurayLib
    
    usd_prim = prim.GetPath()

    inst_name = await mdl_usd.convert_usd_to_mdl(context, usd_prim, path)

    rtncode = (inst_name is not None)

    if rtncode:
        print(f"[omni.mdl.usd_converter] Export to MDL success")
    else:
        print(f"[omni.mdl.usd_converter] Error: Export to MDL failure")

    context.transaction.abort()

    context.set_transaction(None)

    return rtncode


def find_tokens(lines, filter_import_lines = False):
    return mdl_usd.find_tokens(lines, filter_import_lines)
    

# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be
# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
# on_shutdown() is called.
class DiscoveryExtension(omni.ext.IExt): # pragma: no cover
    # ext_id is current extension id. It can be used with extension manager to query additional information, like where
    # this extension is located on filesystem.
    def on_startup(self, ext_id):
        print("[omni.mdl.usd_converter] MDL to USD converter startup")

    def on_shutdown(self):
        print("[omni.mdl.usd_converter]  MDL to USD converter shutdown")
