import os
import sys
import struct


g_pymxsFound   = False
g_MaxPlusFound = False


# pymxs
try:
    from pymxs import runtime as rt
    g_pymxsFound = True
except Exception as e:
    print("pymxs not found")

# MaxPlus
try:
    import MaxPlus
    g_MaxPlusFound = True
except Exception as e:
    print("MaxPlus not found")



def isMaxPlusAvailable():
    return g_MaxPlusFound

def isPymxsAvailable():
    return g_pymxsFound


import nrfile
import nrtools
import nrdump
import nrimp


class kMaterialCreateType:
    Single      = 0
    MultiSubObj = 1


class kConvertMeshTo:
    NoConvert    = 0
    EditablePoly = 1


# kMaterialCreateType.Single
# kMaterialCreateType.MultiSubObj
#g_materialCreateType = kMaterialCreateType.Single

class MaxImportOptions(object):
    def __init__(self):
        self.convertMeshTo      = kConvertMeshTo.NoConvert
        self.materialCreateType = kMaterialCreateType.Single


def writeLog(s):
    print(s)


def matHash(texList):
    res = ""
    for texFile in texList:
        res = res + '__' + os.path.basename(texFile)
    return res


class MaterialManager_MaxPlus(object):
    def __init__(self, options):
        self.options = options
        self.loadedImgs = []
        self.failedImgs = []


    # Material name comparing not work??
    def __isMaterialExist(self, srcMatName):
        #node_name = MaxPlus.WStr(srcMatName)
        mat_lib = MaxPlus.MaterialLibrary.GetSceneMaterialLibrary()
        num_materials = mat_lib.GetNumMaterials()
        #writeLog("SceneMatCnt={:d}".format(num_materials))
        for i in range(0, num_materials):
            mat = mat_lib.GetMaterial(i)
            matName1 = mat.GetFullName()
            #writeLog("Mat={:s}-{:s}".format(matName1, srcMatName))
            if matName1 == MaxPlus.WStr(srcMatName):                      ##### NOT WORK????????????!!
                writeLog("MATERIAL_Found")
                return mat
        return None


    def createMaterial(self, textures, fileDirectory):
        texList = []
        if None != textures:
            texCount = textures.getTexturesCount()
            for i in range(0, texCount):
                texName = fileDirectory + "\\" + textures.getTexture(i).fileName
                texList.append(texName)

        if 0 == len(texList):
            return None


        matName = matHash(texList)
        material = None


        # Need material cache in MaxPlus????
        # Is material library work in MapPlus????
        #material = self.__isMaterialExist(matName)
        #if material:
        #    writeLog("Using existing material")
        #    return material
        #writeLog("Create new material {:s}".format(matName))

        if self.options.maxImpOptions.materialCreateType == kMaterialCreateType.Single:
            # Single
            material = self.__createSingleMaterial(matName, texList)
        elif self.options.maxImpOptions.materialCreateType == kMaterialCreateType.MultiSubObj:
            # MultiSub
            material = self.__createMultiSubMaterial(matName, texList)

        return material


    def __createSingleMaterial(self, matName, texList):
        material = MaxPlus.Factory.CreateDefaultStdMat()
        material.SetName(MaxPlus.WStr(matName))
        ########material.ShowTextureMap = True  ##########???????
        ########material.showInViewport = True
        idx = 0
        for texFile in texList:
            bitmapTex = self.__createBitmapTex(texFile)
            #########bitmapTex.ShowTextureMap = True   ##############????????
            #########bitmapTex.showInViewport = True
            if bitmapTex:
                materialSubMaps = MaxPlus.ISubMap._CastFrom(material)
                #############materialSubMaps.ShowTextureMap = True   ###################?????????
                #############materialSubMaps.showInViewport = True
                materialSubMaps.SetSubTexmap(1 + idx, bitmapTex)
                material.SetEnableMap(1 + idx, True)
                idx = idx + 1
        return material


    def __createMultiSubMaterial(self, matName, texList):
        material = MaxPlus.Factory.CreateDefaultMultiMtl()
        material.SetName(MaxPlus.WStr(matName))
        material.SetNumSubMtls(len(texList))

        idx = 0
        for texFile in texList:
            subMat = material.GetSubMtl(idx)

            if not subMat:
                writeLog("subMat==None")
                continue

            subMat.SetName(MaxPlus.WStr(matName + "_" + str(idx)))

            bitmapTex = self.__createBitmapTex(texFile)
            if not bitmapTex:
                writeLog("bitmapTex==None")
                continue

            materialSubMaps = MaxPlus.ISubMap._CastFrom(subMat)
            if not materialSubMaps:
                writeLog("materialSubMaps==None")
                continue

            materialSubMaps.SetSubTexmap(1, bitmapTex)
            #stdmaterial.SetEnableMap(1, True)

            idx = idx + 1
        return material


    def __createBitmapTex(self, fullpath):
        bitmapTex = None
        try:
            bitmapTex = MaxPlus.Factory.CreateDefaultBitmapTex()
            bitmapTex.SetMapName(fullpath)
            bitmapTex.ReloadBitmapAndUpdate()
            writeLog(u"--> SUCCESSFULL load: {}".format(fullpath))
            self.loadedImgs.append(fullpath)
        except Exception as e:
            writeLog(u"--> FAILED to load: {} Error: {}".format(fullpath, str(e)))
            bitmapTex = None
            self.failedImgs.append(fullpath)
        return bitmapTex


class Importer_MaxPlus(object):
    def __init__(self, options):
        self.options = options
        self.matMgr = MaterialManager_MaxPlus(options)
        return


    def createTexCoords(self, mesh, vert, vertexData, vatrs, faces1, texCoordAttrCompIdxList):
        uvCnt = len(texCoordAttrCompIdxList)
        if 0 == uvCnt:
            writeLog("Texture coords list is empty")
            return

        tex_uvIdx = 0
        for uvIdx in range(0, uvCnt):
            tcAttrComp = texCoordAttrCompIdxList[uvIdx]

            texScatterU = [tcAttrComp.uAttrIdx, tcAttrComp.uCompIdx]  # [attrIdx, compIdx]
            texScatterV = [tcAttrComp.vAttrIdx, tcAttrComp.vCompIdx]  # [attrIdx, compIdx]

            layerTextureCoordinates = nrtools.unpackVertexComponentVaAsList(vert, vertexData, vatrs, [texScatterU, texScatterV])
            if None == layerTextureCoordinates:
                writeLog("UV coordinates not found for U:{:s}-V:{:s}".format(str(texScatterU), str(texScatterV)))
                continue
            
            if 0 == len(layerTextureCoordinates):
                writeLog("No texture coordinates")
                continue
            if len(layerTextureCoordinates[0]) < 2:
                writeLog("UV components < 2")
                continue

            uvwList = []
            for tctc in layerTextureCoordinates:
                uvw = MaxPlus.Point3(float(tctc[0]), 1.0 - float(tctc[1]), 0.0)
                uvwList.append(uvw)

            # Set additional UV in custom map
            customMap = mesh.GetCustomMap(1 + uvIdx)
            if customMap:
                customMap.SetNumTextureVertices(len(uvwList))
                for idx in range(0, len(uvwList)):
                    customMap.SetTextureVertex(idx, uvwList[idx])
                for idx in range(0, len(faces1)):
                    customMap.SetTextureFace(idx, faces1[idx][0], faces1[idx][1], faces1[idx][2])

            if tex_uvIdx == uvIdx:
                mesh.DefaultMap.SetNumTextureVertices(len(uvwList))
                for idx in range(0, len(uvwList)):
                    mesh.DefaultMap.SetTextureVertex(idx, uvwList[idx])   # 1+idx -> ERRROR   VALID: 0-N
                for idx in range(0, len(faces1)):
                    mesh.DefaultMap.SetTextureFace(idx, faces1[idx][0], faces1[idx][1], faces1[idx][2])


    def createMesh(self, vert, vertexData, vatrs,  vert1, vertexData1, vatrs1, meshName, fileDirectory, loadExtraUvData, positions3, indx, textures):
        geom = MaxPlus.Factory.CreateGeomObject(MaxPlus.ClassIds.TriMeshGeometry)
        tri = MaxPlus.TriObject._CastFrom(geom)

        node = MaxPlus.Factory.CreateNode(tri)
        node.SetName(meshName)

        mesh = tri.GetMesh()

        mesh.SetNumVerts(len(positions3))
        #pointList = []
        for idx in range(0, len(positions3)):
            pnt = MaxPlus.Point3(float(positions3[idx][0]), float(positions3[idx][1]), float(positions3[idx][2]))
            #pointList.append(pnt)
            mesh.SetVert(idx, pnt)

        triangles = indx.read()
        faces1 = []
        mesh.SetNumFaces(int(indx.getIndexCount()/3))
        for idx in range(0, int(indx.getIndexCount()/3) ):
            p = struct.unpack_from("iii", triangles, 12*idx)
            #f = MaxPlus.Point3(1 + p[0], 1 + p[1], 1 + p[2])
            faces1.append(p)
            mesh.GetFace(idx).SetVerts(p[0], p[1], p[2])

        mesh.AutoSmooth(45, False)


        # Texcoords
        texCoordAttrList = []
        if nrimp.TextureLoadMode.Auto == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAuto(vatrs)
        elif nrimp.TextureLoadMode.AttributeName == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAttributeName(vatrs, self.options)
        elif nrimp.TextureLoadMode.AttributeIndex == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAttributeIndex(vatrs, self.options)

        if not loadExtraUvData:
            self.createTexCoords(mesh, vert, vertexData, vatrs, faces1, texCoordAttrList)
        else:
            self.createTexCoords(mesh, vert1, vertexData1, vatrs1, faces1, texCoordAttrList)


        # Create material if textures presented
        material = self.matMgr.createMaterial(textures, fileDirectory)
        if material:
            node.Material = material
            # Show Shaded Material in Viewport
            MaxPlus.Core.EvalMAXScript('showTextureMap $'+node.Name + '.material  true') 


        #mod = MaxPlus.Factory.CreateObjectModifier(MaxPlus.ClassIds.Edit_Mesh)
        ##mod = MaxPlus.Factory.CreateObjectModifier(MaxPlus.ClassIds.Edit_Patch)  
        #node.AddModifier(mod)
        #node.Collapse(True)
        #node.Convert(MaxPlus.ClassIds.PolyMeshObject)

        if self.options.maxImpOptions.convertMeshTo == kConvertMeshTo.EditablePoly:
            mod = MaxPlus.Factory.CreateObjectModifier(MaxPlus.ClassIds.Edit_Patch)  
            node.AddModifier(mod)
            node.Collapse(True)
            node.Convert(MaxPlus.ClassIds.PolyMeshObject)


        mesh.InvalidateGeomCache()
        mesh.InvalidateTopologyCache()

        return True


#########################################################################


class MaterialManager_pymxs(object):    
    def __init__(self, options):
        self.options = options
        self.loadedImgs = []
        self.failedImgs = []


    def __isMaterialExist(self, matName):
        for mat in rt.sceneMaterials:
            if mat.name == matName:
                return mat
        return None


    def createMaterial(self, textures, fileDirectory):
        texList = []
        if None != textures:
            texCount = textures.getTexturesCount()
            for i in range(0, texCount):
                texName = fileDirectory + "\\" + textures.getTexture(i).fileName
                texList.append(texName)

        if 0 == len(texList):
            return None


        matName = matHash(texList)

        material = self.__isMaterialExist(matName)
        if material:
            #writeLog("Mat exist {:s}".format(matName))
            return material

        writeLog(u"Creating new material {}".format(matName))

        if self.options.maxImpOptions.materialCreateType == kMaterialCreateType.Single:
            # Single
            material = self.__createSingleMaterial(matName, texList)
        elif self.options.maxImpOptions.materialCreateType == kMaterialCreateType.MultiSubObj:
            # MultiSub
            material = self.__createMultiSubMaterial(matName, texList)
        return material


    def __createBitmap(self, fullpath):
        bitmap = None
        try:
            bitmap = rt.openBitMap(fullpath)
            writeLog(u"--> SUCCESSFULL load: {}".format(fullpath))
            self.loadedImgs.append(fullpath)
        except Exception as e:
            writeLog(u"--> FAILED to load: {}".format(fullpath))
            bitmap = None
            self.failedImgs.append(fullpath)
        return bitmap


    def __createSingleMaterial(self, matName, texList):
        idx = 0
        #material = rt.StandardMaterial(name=matName)
        material = rt.PhysicalMaterial(name=matName)
        rt.showTextureMap(material, True)
        for texFile in texList:
            bitmap = self.__createBitmap(texFile)
            if bitmap:
                bmt = rt.bitmapTex()
                bmt.bitmap = bitmap
                rt.SetSubTexmap(material, 2 + idx, bmt)  # Set BASE COLOR MAP
            idx = idx + 1
        return material


    def __createMultiSubMaterial(self, matName, texList):
        idx = 0
        material = rt.Multimaterial(name=matName, numsubs=len(texList))
        for texFile in texList:
            subMat = rt.getSubMtl(material, 1 + idx)
            bitmap = self.__createBitmap(texFile)
            if bitmap:
                bmt = rt.bitmapTex()
                bmt.bitmap = bitmap
                rt.SetSubTexmap(subMat, 2, bmt)  # Set BASE COLOR MAP
            idx = idx + 1
        return material


class Importer_pymxs(object):
    def __init__(self, options):
        self.options = options
        self.matMgr = MaterialManager_pymxs(options)
        return


    def createTexCoords(self, mesh, vert, vertexData, vatrs, faces1, texCoordAttrCompIdxList):
        uvCnt = len(texCoordAttrCompIdxList)
        if 0 == uvCnt:
            writeLog("Texture coords list is empty")
            return

        tex_uvIdx = 0
        rt.meshop.setNumMaps(mesh, 1 + uvCnt)
        for uvIdx in range(0, uvCnt):
            tcAttrComp = texCoordAttrCompIdxList[uvIdx]

            texScatterU = [tcAttrComp.uAttrIdx, tcAttrComp.uCompIdx]  # [attrIdx, compIdx]
            texScatterV = [tcAttrComp.vAttrIdx, tcAttrComp.vCompIdx]  # [attrIdx, compIdx]

            layerTextureCoordinates = nrtools.unpackVertexComponentVaAsList(vert, vertexData, vatrs, [texScatterU, texScatterV])
            if None == layerTextureCoordinates:
                writeLog("UV coordinates not found for U:{:s}-V:{:s}".format(str(texScatterU), str(texScatterV)))
                continue
            
            if 0 == len(layerTextureCoordinates):
                writeLog("No texture coordinates")
                continue
            if len(layerTextureCoordinates[0]) < 2:
                writeLog("UV components < 2")
                continue

            uvwList = []

            for tctc in layerTextureCoordinates:
                uvw = rt.Point3(float(tctc[0]), 1.0 - float(tctc[1]), 0.0)
                uvwList.append(uvw)

            rt.meshop.setMapSupport(mesh, 1 + uvIdx, True);

            # Assign Default Texture mapping
            if uvIdx == tex_uvIdx:
                rt.buildTVFaces(mesh, False)
                for idx in range(0, len(faces1)):
                    rt.setTVFace(mesh, 1 + idx, faces1[idx])

                rt.meshop.setNumTVerts(mesh, len(uvwList))
                for idx in range(0, len(uvwList)):
                    rt.setTVert(mesh, 1 + idx, uvwList[idx])   # 3ds Max 2018 bug

            # Set UV-channel data
            for idx in range(0, len(uvwList)):
                rt.meshop.setMapVert(mesh, 1 + uvIdx, 1 + idx,  uvwList[idx])


    def createMesh(self, vert, vertexData, vatrs,  vert1, vertexData1, vatrs1, meshName, fileDirectory, loadExtraUvData, positions3, indx, textures):
        pointList = []
        for idx in range(0, len(positions3)):
            pnt = rt.Point3(float(positions3[idx][0]), float(positions3[idx][1]), float(positions3[idx][2]))
            pointList.append(pnt)

        # Indexes started from 1 ?!
        triangles = indx.read()
        faces1 = []
        for idx in range(0, int(indx.getIndexCount()/3) ):
            p = struct.unpack_from("iii", triangles, 12*idx)
            f = rt.Point3(1 + p[0], 1 + p[1], 1 + p[2])
            faces1.append(f)

        mesh = None

        # Create material from textures
        material1 = self.matMgr.createMaterial(textures, fileDirectory)
        if material1:
            mesh = rt.mesh(name=meshName, vertices=pointList, faces=faces1, material=material1)
        else:
            mesh = rt.mesh(name=meshName, vertices=pointList, faces=faces1)

        # Texcoords
        texCoordAttrList = []
        if nrimp.TextureLoadMode.Auto == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAuto(vatrs)
        elif nrimp.TextureLoadMode.AttributeName == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAttributeName(vatrs, self.options)
        elif nrimp.TextureLoadMode.AttributeIndex == self.options.textureLoadMode:
            texCoordAttrList = nrimp.createTexCoordAttributeIndex(vatrs, self.options)

        if not loadExtraUvData:
            self.createTexCoords(mesh, vert, vertexData, vatrs, faces1, texCoordAttrList)
        else:
            self.createTexCoords(mesh, vert1, vertexData1, vatrs1, faces1, texCoordAttrList)

        rt.update(mesh)
        return True


#########################################################################


class MaxImporter(object):
    def __init__(self, options):
        self.options = options
        self.totalFilesCount = 0
        self.totalCreated = 0
        self.totalSkipped = 0

        self.postVsImportFlag = False
        if  nrimp.GeometryLoadMode.PostVs == self.options.geometryLoadMode:
            self.postVsImportFlag = True

        self.imp_MaxPlus = Importer_MaxPlus(options)
        self.imp_pymxs   = Importer_pymxs(options)

        self.loadExtraVertexData = self.options.loadExtraUvData

        self._maxNrSize = 0
        self._maxMeshName = ''


    def importTriangleMesh(self, nrmesh, indx, vert, vertexData, vatrs, positions3, meshName, fileDirectory):
        # Extra vatrs/vertexes
        loadExtraUvData = False
        vert1  = None
        vatrs1 = None
        vertexData1 = None
        if self.loadExtraVertexData:
            while (True):
                vert1  = nrmesh.getVertexes(1)
                vatrs1 = nrmesh.getVertexAttributes(1)
                if (not vert1) or (not vatrs1):
                    break

                vertexData1 = vert1.read()
                if not vertexData1:
                    break

                if vert.getVertexCount() != vert1.getVertexCount():
                    writeLog("Vertex count != Extra vertex count {:d} != {:d}".format(vert.getVertexCount(), vert1.getVertexCount()))
                    break

                loadExtraUvData = True
                break

        # Get Textures object (TXTR tag in rip file)
        #  textures = texture files list
        textures = nrmesh.getTextures()

        # Create mesh pymxs/MaxPlus
        if isMaxPlusAvailable():
            meshCreated = self.imp_MaxPlus.createMesh(vert, vertexData, vatrs,  vert1, vertexData1, vatrs1, meshName, fileDirectory, loadExtraUvData, positions3, indx, textures)
        elif isPymxsAvailable():
            meshCreated = self.imp_pymxs.createMesh(vert, vertexData, vatrs,  vert1, vertexData1, vatrs1, meshName, fileDirectory, loadExtraUvData, positions3, indx, textures)
        else:
            writeLog("Import API Not Found pymxs/MaxPlus")


    def importMesh(self, filename):
        res = False
        try:
            res = self.importMeshImpl(filename)
        except Exception as e:
            writeLog("Import exception: {}".format(str(e)))
        return res


    def importMeshImpl(self, filename):
        nr = nrfile.NRFile()
        if not nr.parse(filename):
            writeLog("Ninja Ripper file parsing failed: {}".format(nr.getErrorString()))
            return False

        #print("FileNameClass={}".format(filename.__class__.__name__))
        writeLog(u"Loading: {}".format(filename))

        fileDirectory = os.path.dirname(os.path.abspath(filename))
        
        createdMeshsCount = 0

        for meshIdx in range(0, nr.getMeshCount()):
            nrmesh = nr.getMesh(meshIdx)

            # ---Fast checks for mesh stage---
            if self.postVsImportFlag:
                # PostVS processing
                if nrfile.ShaderStage.PreVs == nrmesh.getShaderStage():
                    writeLog("Skipping PreVs")
                    continue
            else:
                # PreVS processing
                if nrfile.ShaderStage.PreVs != nrmesh.getShaderStage():
                    writeLog("Skipping PostVs")
                    continue

            vatrs = nrmesh.getVertexAttributes(0)
            if None == vatrs:
                writeLog("vertexAttribs == Null")
                continue

            if self.postVsImportFlag:
                # PostVS attribute with idx==0 is always POSITION with four components (homogeneous coordinates xyzw)
                if vatrs.getAttr(0).compCount != 4:
                    writeLog("Post vs position not 4 component")
                    continue


            indx = nrmesh.getIndexes(0)
            if None == indx:
                writeLog("indexes == Null")
                continue

            vert = nrmesh.getVertexes(0)
            if None == vert:
                writeLog("vertexes == Null")
                continue

            vertexData = vert.read()
            if None == vertexData:
                writeLog("vertexData == Null")
                continue

            positions3 = None
            if self.postVsImportFlag:
                # PostVS
                positions3 = nrtools.restorePositionAsList(vert, vertexData, vatrs.getAttr(0), self.options.projMat)
                if None == positions3:
                    writeLog("positions3 == Null")
                    continue
            else:
                # PreVS
                if nrimp.GeometryLoadMode.PreVs_Auto == self.options.geometryLoadMode:
                    posAttr = nrimp.createPreVsAuto(vatrs)
                    positions3 = nrtools.unpackVertexComponentVaAsList(vert, vertexData, vatrs, [posAttr[0], posAttr[1], posAttr[2]])
                elif nrimp.GeometryLoadMode.PreVs_Manual == self.options.geometryLoadMode:
                    positions3 = nrtools.unpackVertexComponentVaAsList(vert, vertexData, vatrs, [[self.options.prevs_xAttrIdx, self.options.prevs_xCompIdx], [self.options.prevs_yAttrIdx, self.options.prevs_yCompIdx], [self.options.prevs_zAttrIdx, self.options.prevs_zCompIdx]] )
                if None == positions3:
                    writeLog("positions3 == Null")
                    continue

            meshName = os.path.basename(filename)
            meshName = os.path.splitext(meshName)[0]
            meshCreated = False
            if nrfile.IndexTopology.TriangleList == indx.getIndexTopology():
                # Create triangle mesh
                meshCreated = self.importTriangleMesh(nrmesh, indx, vert, vertexData, vatrs, positions3, meshName, fileDirectory)
            else:
                # Create line/point mesh
                writeLog("Import not realized for index topology={:s}".format(nrdump.topologyToStr(indx.getIndexTopology())))
            

            if meshCreated:
                createdMeshsCount = createdMeshsCount + 1

        self.totalCreated    = self.totalCreated + createdMeshsCount
        self.totalSkipped    = self.totalSkipped + (nr.getMeshCount() - createdMeshsCount)
        self.totalFilesCount = self.totalFilesCount + 1
        return True
