__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
import platform
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 Vector3D(Structure):
    _fields_ = [ ("x", c_float),("y", c_float),("z", c_float), ]
if '32bit' in platform.architecture():
    class MaterialProperty(Structure):
        _fields_=[("mKey", String),("mSemantic", c_uint),("mIndex", c_uint),("mDataLength", c_uint),("mType", c_uint),("mData",POINTER(c_char)),]
else: #64bit has extra padding :(
    class String4(Structure):
        _fields_ = [ ("length", c_uint),("pad",c_uint),("data", c_char*1024), ]
    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 = ".",'.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)
            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, ROTATE_X90 ):
    t = time.time()
    model = dll.aiImportFile( path, 2 ).contents #aiProcess_JoinIdenticalVertices=2
    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 in [2,7]: material.add_texture(texture, "UV", "SPECULARITY" ) #7="HARDNESS"
                #elif p.mSemantic in [3,4]: material.add_texture(texture, "UV", "AMBIENT" ) #4="EMIT"
                elif p.mSemantic in [5,6]:
                    material.add_texture(texture, "UV", "NORMAL" ) #5="HEIGHT"
                    texture.normal_map = True
                    texture.normal_space = "TANGENT"
                elif p.mSemantic == 8:
                    material.add_texture(texture, "UV", "ALPHA" )
                    texture.use_alpha = False
                    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)
        if ROTATE_X90:
            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) ] )
        else:
            me.verts.foreach_set("co", [f for i in range(mesh.mNumVertices) for f in (mesh.mVertices[i].x, mesh.mVertices[i].y, mesh.mVertices[i].z) ] )
        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.double_sided = False
        me.vertex_normal_flip = False
        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.add_material(materials[mesh.mMaterialIndex])
        me.update()
        ob = bpy.data.objects.new("Mesh", me)
        context.scene.objects.link(ob)
        context.scene.objects.active = ob
        #bpy.ops.object.select_all(action='DESELECT')
        bpy.ops.object.mode_set(mode='EDIT')
        #bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.mesh.faces_shade_smooth()
        bpy.ops.object.mode_set(mode='OBJECT')
        #bpy.ops.object.shade_smooth()
    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= "")
    ROTATE_X90 = BoolProperty(name="-X90", description="Rotate X 90.", default= True)

    def execute(self, context):
        load_assimp(self.properties.path, context, self.properties.ROTATE_X90 )
        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()