import_m3.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####


import bpy, mathutils, time, struct, os.path
from bpy.props import *

##################
## Struct setup ##
##################

verFlag = False        #Version flag (MD34 == True, MD33 == False)

class ref:
    fmt = 'LL'
    def __init__(self, file):
        global verFlag
        if verFlag: self.fmt += 'L'    #Extra unknown...
        _s=file.read(struct.calcsize(self.fmt))
        self.entries, self.refid = struct.unpack(self.fmt, _s)[:2]
    @classmethod
    def size(cls):
        global verFlag
        return struct.calcsize(cls.fmt + ('L' if verFlag else ''))


class animref:
    fmt = 'HHL'
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        self.flags, self.animflags, self.animid = struct.unpack(self.fmt, _s)


class Tag:
    fmt = '4sLLL'
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        self.name, self.ofs, self.nTag, self.version = struct.unpack(self.fmt, _s)


class matrix:
    fmt='f'*16
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        self.mat = struct.unpack(self.fmt, _s)


class vect:
    fmt='fff'
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        self.v = struct.unpack(self.fmt, _s)


class vertex:
    fmt="4B4b4B%dH4B"
    ver = { 0x020000: 2, 0x060000: 4, 0x0A0000: 6, 0x120000: 8 }
    def __init__(self, file, flag):
        self.pos = vect(file)
        fmt = self.fmt % (self.ver[flag])
        _s=file.read(struct.calcsize(fmt))
        print(file.tell())
        _s = struct.unpack(fmt, _s)
        self.boneWeight = _s[0:4]
        self.boneIndex = _s[4:8]
        self.normal = _s[8:12]
        self.uv = _s[12:14]
        self.tan = _s[-4:]    #Skipping the middle ukn value if needed
        self.boneWeight = [b/255.0 for b in self.boneWeight]
        self.normal = [x*2.0/255.0-1.0 for x in self.normal]
        self.tan = [x*2.0/255.0-1.0 for x in self.tan]
        self.uv = [x/2046.0 for x in self.uv]
        self.uv[1] = 1.0 - self.uv[1]
    @classmethod
    def size(cls, flag):
        return struct.calcsize(cls.fmt % (cls.ver[flag]))


class quat:
    fmt='ffff'
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        self.v = struct.unpack(self.fmt, _s)
        self.v = [self.v[-1], self.v[0], self.v[1], self.v[2]]    #Quats are stored x,y,z,w - this fixes it


class bone:
    def __init__(self, file):
        file.read(4)    #ukn1
        self.name = ref(file)
        self.flag, self.parent, _ = struct.unpack('LhH',file.read(8))
        self.posid = animref(file)
        self.pos = vect(file)
        file.read(4*4)    #ukn
        self.rotid = animref(file)
        self.rot = quat(file)
        file.read(4*5)    #ukn
        self.scaleid = animref(file)
        self.scale = vect(file)
        vect(file)        #ukn
        file.read(4*6)    #ukn


class div:
    def __init__(self, file):
        self.faces = ref(file)
        self.regn = ref(file)
        self.bat = ref(file)
        self.msec = ref(file)
        file.read(4)    #ukn


class regn:
    fmt = 'LHHLL6H'
    def __init__(self, file):
        _s=file.read(struct.calcsize(self.fmt))
        _ukn1, self.ofsVert, self.nVerts, self.ofsIndex, self.nIndex, \
            self.boneCount, self.indBone, self.nBone = struct.unpack(self.fmt, _s)[:8]


class mat:
    def __init__(self, file):
        self.name = ref(file)
        file.read(4*10)    #ukn
        self.layers = [ref(file) for _ in range(13)]
        file.read(4*15)    #ukn


class layr:
    def __init__(self, file):
        file.read(4)
        self.name = ref(file)
        #Rest not implemented.


class hdr:
    fmt = '4sLL'
    def __init__(self, file):
        _s = file.read(struct.calcsize(self.fmt))
        self.magic, self.ofsTag, self.nTag = struct.unpack(self.fmt, _s)
        self.MODLref = ref(file)


class MODL:
    def __init__(self, file, flag=20):
        global verFlag
        self.name = ref(file)
        self.ver = struct.unpack('L',file.read(4))[0]
        self.seqHdr = ref(file)
        self.seqData = ref(file)
        self.seqLookup = ref(file)
        file.read(0x1C if verFlag else 0x14)            #ukn1
        self.bones = ref(file)
        file.read(4)            #ukn2
        self.flags = struct.unpack('L',file.read(4))[0]
        self.vert = ref(file)
        self.views = ref(file)
        self.boneLookup = ref(file)
        self.extents = [vect(file), vect(file)]
        self.radius = struct.unpack('f',file.read(4))[0]
        if verFlag: file.read(4)            #ukn MD34 addition
        if not verFlag:
            if flag == 20: file.read(0x2C)
            else: file.read(0x34)
        else:
            if flag == 20: file.read(0x30)
            else: file.read(0x3C)
        self.attach = ref(file)
        file.read(5*ref.size())
        self.materialsLookup = ref(file)
        self.materials = ref(file)
        file.read(ref.size())
        if not verFlag:
            file.read(0x90)
        else: file.read(0xD8)
        self.iref = ref(file)


def read(file, context, op):
    '''Imports as an m3 file'''
    
    global verFlag
    h=hdr(file)
    if h.magic == b'43DM':
        print('m3_import: !WARNING! MD34 files not full tested...')
        verFlag = True
    elif h.magic == b'33DM':
        verFlag = False
    else:
        raise Exception('m3_import: !ERROR! Not a valid or supported m3 file')
    file.seek(h.ofsTag)    #Jump to the Tag table
    print('m3_import: !INFO! Reading TagTable...')
    tagTable = [Tag(file) for _ in range(h.nTag)]
    file.seek(tagTable[h.MODLref.refid].ofs)
    m = MODL(file, tagTable[h.MODLref.refid].version)
    if not m.flags & 0x20000:
        raise Exception('m3_import: !ERROR! Model doesn\'t contain any vertices')
    vert_flags = m.flags & 0x1E0000        #Mask out the vertex version
    file.seek(tagTable[m.vert.refid].ofs)
    print('m3_import: !INFO! Reading Vertices...')
    verts = [ vertex(file, vert_flags) for _ in range(int(tagTable[m.vert.refid].nTag / vertex.size(vert_flags))) ]
    file.seek(tagTable[m.views.refid].ofs)
    d = div(file)
    file.seek(tagTable[d.faces.refid].ofs)
    print('m3_import: !INFO! Reading Faces...')
    rawfaceTable = struct.unpack('H'*(tagTable[d.faces.refid].nTag), file.read(tagTable[d.faces.refid].nTag*2))
    faceTable = []
    for i in range(1, len(rawfaceTable)+1):
        faceTable.append(rawfaceTable[i-1])
        if i % 3 == 0: faceTable.append(0)    #Add a zero for the fourth index to the face.
    print('m3_import: !INFO! Adding Geometry...')
    mesh = bpy.data.meshes.new(os.path.basename(op.properties.filepath))
    mobj = bpy.data.objects.new(os.path.basename(op.properties.filepath), mesh)
    context.scene.objects.link(mobj)
    v = []
    for vert in verts:        #"Flatten" the vertex array...
        v.extend(vert.pos.v)
    mesh.vertices.add(len(verts))
    mesh.faces.add(len(rawfaceTable)//3)
    mesh.vertices.foreach_set('co', v)
    mesh.faces.foreach_set('vertices_raw', faceTable)
    uvtex = mesh.uv_textures.new()
    for i, face in enumerate(mesh.faces):
        uf = uvtex.data[i]
        uf.uv1 = verts[faceTable[i*4+0]].uv
        uf.uv2 = verts[faceTable[i*4+1]].uv
        uf.uv3 = verts[faceTable[i*4+2]].uv
        uf.uv4 = (0,0)
    print('m3_import: !INFO! Importing materials...')
    material = bpy.data.materials.new('Mat00')
    mesh.materials.append(material)
    file.seek(tagTable[m.materials.refid].ofs)
    mm = mat(file)
    for map, i in [('use_map_diffuse', 0), ('use_map_specular', 2), ('use_map_normal', 9)]:
        file.seek(tagTable[mm.layers[i].refid].ofs)
        nref = layr(file).name
        file.seek(tagTable[nref.refid].ofs)
        name = bytes.decode(file.read(nref.entries-1))
        path = os.path.join(os.path.dirname(op.properties.filepath), os.path.basename(str(name)))
        tex = bpy.data.textures.new(name=os.path.basename(path), type='IMAGE')
        if os.path.exists(path):
            tex.image = bpy.data.images.load(path)
            print("m3_import: !INFO! Loaded %s" % (path))
        else:
            print("m3_import: !WARNING! Cannot find texture \"%s\"" % (path))
        mtex = material.texture_slots.add()
        mtex.texture = tex
        mtex.texture_coords = 'UV'
        mtex.use_map_color_diffuse = (i == 0)
        setattr(mtex, map, True)    # <- This is really dumb to have to do, just expose a bitmask or something.


class M3Importer(bpy.types.Operator):
    '''Import from M3 file format (.m3)'''
    
    bl_idname = "import_mesh.blizzard_m3"
    bl_label = 'Import M3'

    # List of operator properties, the attributes will be assigned
    # to the class instance from the operator settings before calling.

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

    def execute(self, context):
        t = time.clock()
        with open(self.properties.filepath, 'rb') as file:
            print('Importing file', self.properties.filepath)
            read(file, context, self)
        print('Finished importing in', time.clock() - t, 'seconds')
        return {'FINISHED'}

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


def menu_func(self, context):
    self.layout.operator(M3Importer.bl_idname, text="M3 (.m3)")


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


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


if __name__ == "__main__":
    register()