__author__= "Matthias Fauconneau"
__url__= ['blenderartists.org']
__version__= "0.1"
__bpydoc__= """This script imports any files supported by Assimp to Blender."""

import time
import ctypes
import os
from ctypes import POINTER, Structure, c_char, c_char_p, c_uint, c_int, c_float, cast, sizeof, pointer
import bpy
from bpy.props import *

class Face(Structure):
    _fields_ = [ ("mNumIndices", c_uint), ("mIndices", POINTER(c_uint)), ]
class String(Structure):
    _fields_ = [ ("length", c_uint),("data", c_char*1024), ]
class String4(Structure):
    _fields_ = [ ("length", c_uint),("pad",c_uint),("data", c_char*1024), ]
class Vector3D(Structure):
    _fields_ = [ ("x", c_float),("y", c_float),("z", c_float), ]
class MaterialProperty(Structure):
    _fields_=[("mKey", String4),("mSemantic", c_uint),("mIndex", c_uint),("mDataLength", c_uint),("mType", c_uint),("mData",POINTER(c_char)),]
class Material(Structure):
    _fields_ = [("mProperties", POINTER(POINTER(MaterialProperty))),("mNumProperties", c_uint),("mNumAllocated", c_uint),]
class Mesh(Structure):
    _fields_ = [("mPrimitiveTypes", c_uint),("mNumVertices", c_uint),("mNumFaces", c_uint),
        ("mVertices", POINTER(Vector3D)),("mNormals", POINTER(Vector3D)),("mTangents", POINTER(Vector3D)),("mBitangents", POINTER(Vector3D)),
        ("mColors", POINTER(c_uint)*4),("mTextureCoords", POINTER(Vector3D)*4),("mNumUVComponents", c_uint*4),("mFaces", POINTER(Face)),
        ("mNumBones", c_uint),("mBones", POINTER(POINTER(c_uint))),("mMaterialIndex", c_uint),]
class Scene(Structure):
    _fields_ = [("mFlags", c_uint),("mRootNode", POINTER(c_uint)),("mNumMeshes", c_uint),("mMeshes", POINTER(POINTER(Mesh))),
    ("mNumMaterials", c_uint),("mMaterials", POINTER(POINTER(Material))),("mNumAnimations", c_uint),("mAnimations", POINTER(POINTER(c_uint))),
    ("mNumTextures", c_uint),("mTextures", POINTER(POINTER(c_uint))),("mNumLights", c_uint),("mLights", POINTER(POINTER(c_uint))),
    ("mNumCameras", c_uint),("mCameras", POINTER(POINTER(c_uint))),]

def search_library():
    # silence 'DLL not found' message boxes on win
    try: ctypes.windll.kernel32.SetErrorMode(0x8007)
    except AttributeError: pass

    if os.name=='posix':
        folder,ext = '/usr/local/lib/','.so'
    elif os.name=='nt':
        folder,ext = os.path.dirname(__file__),'.dll'
    for filename in os.listdir(folder):
            if filename.lower().find('assimp') > 0 and os.path.splitext(filename)[-1].lower() == ext:
                library = os.path.join(folder, filename)
                print( 'Try ',library )
                try: dll = ctypes.cdll.LoadLibrary(library)
                except: continue
                try: load = dll.aiImportFile
                except AttributeError: pass
                else:
                    dll.aiImportFile.restype = POINTER(Scene)
                    dll.aiImportFile.argtypes = [c_char_p, c_uint]
                    return dll;
    return None

dll = search_library()

# limited replacement for BPyImage.comprehensiveImageLoad
def load_image(imagepath, dirname):
        if os.path.exists(imagepath): return bpy.data.images.load(imagepath)
        variants = [os.path.join(dirname, imagepath), os.path.join(dirname, os.path.basename(imagepath))]
        for path in variants:
                if os.path.exists(path): return bpy.data.images.load(path)
                else: print(path, "doesn't exist")
        return None

def load_assimp( path, context ):
        t = time.time()
        model = dll.aiImportFile( path, 0 ).contents
        materials = []
        for index in range(model.mNumMaterials):
                for i in range(model.mMaterials[index].contents.mNumProperties):
                        p = model.mMaterials[index].contents.mProperties[i].contents
                        key = p.mKey.data
                        if p.mType == 1: value = [cast(p.mData, POINTER(c_float) )[i] for i in range(p.mDataLength//sizeof(c_float))]
                        elif p.mType == 3: value = cast(p.mData, POINTER(String)).contents.data
                        elif p.mType == 4: value = [cast(p.mData, POINTER(c_int) )[i] for i in range(p.mDataLength//sizeof(c_int))]
                        else: value = p.mData[:p.mDataLength]
                        print( "    %s: %s" % (key, value) )
                        if( key == "?mat.name" ):
                                material = bpy.data.materials.new( value )
                        if( key == "$clr.diffuse" ):
                                material.diffuse_color = ( value[0], value[1], value[2] )
                        if( key == "$clr.specular" ):
                                material.specular_color = ( value[0], value[1], value[2] )
                        if( key == "$mat.shininess" ):
                                material.specular_hardness = int(value[0])
                        if key == "$tex.file":
                                texture= bpy.data.textures.new(value)
                                texture.type= 'IMAGE'
                                texture = texture.recast_type()
                                image = load_image( value, os.path.dirname(path) )
                                has_data = image.has_data if image else False
                                if image: texture.image=image
                                if p.mSemantic == 1 and has_data and image.depth == 32:
                                        material.add_texture(texture, "UV", ("COLOR", "ALPHA"))
                                        texture.mipmap = True
                                        texture.interpolation = True
                                        texture.use_alpha = True
                                        material.alpha = 0.0
                                elif p.mSemantic == 1: material.add_texture(texture, "UV", "COLOR" )
                                elif p.mSemantic == 2: material.add_texture(texture, "UV", "SPECULARITY" )
                                elif p.mSemantic == 3: material.add_texture(texture, "UV", "AMBIENT" )
                                elif p.mSemantic == 4: material.add_texture(texture, "UV", "EMIT" )
                                elif p.mSemantic == 5: material.add_texture(texture, "UV", "NORMAL" ) #"HEIGHT"
                                elif p.mSemantic == 6: material.add_texture(texture, "UV", "NORMAL" )
                                elif p.mSemantic == 7: material.add_texture(texture, "UV", "HARDNESS") #"SHININESS"
                                elif p.mSemantic == 8:
                                        material.add_texture(texture, "UV", "ALPHA" )
                                        material.transparency = True
                                        material.transparency_method = "Z_TRANSPARENCY"
                                        material.alpha = 0.0
                materials.append(material)

        for index in range(model.mNumMeshes):
                mesh = model.mMeshes[index].contents
                me = bpy.data.meshes.new("Mesh")
                me.add_geometry(mesh.mNumVertices, 0, mesh.mNumFaces)
                me.verts.foreach_set("co", [f for i in range(mesh.mNumVertices) for f in (mesh.mVertices[i].x, mesh.mVertices[i].z, mesh.mVertices[i].y) ] )
                def treat_face(f):
                        if f.mNumIndices == 3: return f.mIndices[2], f.mIndices[1], f.mIndices[0], 0
                        return  f.mIndices[2], f.mIndices[1], f.mIndices[0], f.mIndices[3]
                me.faces.foreach_set("verts_raw", [ f for i in range(mesh.mNumFaces) for f in treat_face(mesh.mFaces[i]) if mesh.mFaces[i].mNumIndices >= 3 ] )
                me.add_uv_texture()
                for i in range(mesh.mNumFaces):
                        if mesh.mFaces[i].mNumIndices >= 3:
                                tface = me.uv_textures[0].data[i]
                                tface.uv1= mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[2]].x, mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[2]].y
                                tface.uv2= mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[1]].x, mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[1]].y
                                tface.uv3= mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[0]].x, mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[0]].y
                                if mesh.mFaces[i].mNumIndices == 4:
                                        tface.uv4= mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[3]].x, mesh.mTextureCoords[0][mesh.mFaces[i].mIndices[3]].y
                me.update()
                ob = bpy.data.objects.new("Mesh", me)
                context.scene.objects.link(ob)
                me.add_material(materials[mesh.mMaterialIndex])

        dll.aiReleaseImport(pointer(model))
        print( "%.3f" % (time.time()-t) )

class IMPORT_OT_assimp(bpy.types.Operator):
        '''Load any file supported by Assimp (.obj .3ds .ase .lwo .md3 .ply ...)'''
        bl_idname = "import_scene.assimp"
        bl_label = "Import using Assimp"

        path = StringProperty(name="File Path", description="File path used for importing the file", maxlen= 1024, default= "")

        def execute(self, context):
                load_assimp(self.properties.path, context)
                return {'FINISHED'}

        def invoke(self, context, event):
                wm = context.manager
                wm.add_fileselect(self)
                return {'RUNNING_MODAL'}

menu_func = lambda self, context: self.layout.operator(IMPORT_OT_assimp.bl_idname, text="using Assimp (.obj .3ds .ase .lwo .md3 .ply ...)")

def register():
        bpy.types.register(IMPORT_OT_assimp)
        bpy.types.INFO_MT_file_import.append(menu_func)

def unregister():
        bpy.types.unregister(IMPORT_OT_assimp)
        bpy.types.INFO_MT_file_import.remove(menu_func)

if __name__ == "__main__":
        register()
