# bitcpy
# CaveTroll

bl_info = {
    "name": "Legend of Grimrock Mesh Format (.mesh)",
    "author": "",
    "version": (1, 0, 1),
    "blender": (2, 5, 7),
    "api": 36339,
    "location": "File > Import > Legend of Grimrock Mesh (.mesh)",
    "description": "Import Legend of Grimrock Meshes (.mesh)",
    "warning": "",
    "wiki_url": "",
    "tracker_url": "",
    "category": "Import-Export"}

import os
import struct

import bpy
from mathutils import *

from bpy.props import *
from bpy_extras.io_utils import ExportHelper, ImportHelper
from bpy_extras.image_utils import load_image

class _mesh_material(object):
    __slots__ = (
        "bl_mat",
        "bl_image",
        "name",
        "index_offset",
        "num_primitives",
        )

    def __init__(self):
        self.bl_mat = None
        self.bl_image = None
        self.name = "<Empty>"
        self.index_offset = 0
        self.num_primitives = 0

# Reads file magic from file
def read_magic(file_object, endian = '<'):
    data = struct.unpack(endian+'4s', file.read(4))[0]
    return data;

# Read signed integer from file
def read_int(file_object, endian = '<'):
    data = struct.unpack(endian+'i', file_object.read(4))[0]
    return data

def read_int2(file_object, endian = '<'):
    data = struct.unpack(endian+'ii', file_object.read(8))
    return data

def read_int3(file_object, endian = '<'):
    data = struct.unpack(endian+'iii', file_object.read(12))
    return data

def read_int4(file_object, endian = '<'):
    data = struct.unpack(endian+'iiii', file_object.read(16))
    return data

# Read floating point number from file
def read_float(file_object, endian = '<'):
    data = struct.unpack(endian+'f', file_object.read(4))[0]
    return data

def read_float2(file_object, endian = '<'):
    data = struct.unpack(endian+'ff', file_object.read(8))
    return data

def read_float3(file_object, endian = '<'):
    data = struct.unpack(endian+'fff', file_object.read(12))
    return data

def read_float4(file_object, endian = '<'):
    data = struct.unpack(endian+'ffff', file_object.read(16))
    return data

# Read unsigned bytes from file
def read_byte(file_object, endian = '<'):
    data = struct.unpack(endian+'B', file_object.read(1))[0]
    return data

def read_byte2(file_object, endian = '<'):
    data = struct.unpack(endian+'BB', file_object.read(2))
    return data

def read_byte3(file_object, endian = '<'):
    data = struct.unpack(endian+'BBB', file_object.read(3))
    return data

def read_byte4(file_object, endian = '<'):
    data = struct.unpack(endian+'BBBB', file_object.read(4))
    return data

# Read string from file
def read_string(file_object, num, endian = '<'):
    raw_string = struct.unpack(endian+str(num)+'s', file_object.read(num))[0]
    data = raw_string.decode("utf-8", "ignore")
    return data


# Build content base path
def build_assets_path(filename):
    #filename = "C:/grimrock/assets/models/wall_sets/dungeon/floor_01.mesh"
    pathname = os.path.dirname(filename)

    index = pathname.rfind("assets")
    if index < 0:
        return pathname

    return pathname[0:index]


# Parse the '/assets/materials/default.materials' file
def parse_material_file( path_assets, find_material ):

    # diffuse, normal, specular
    texture_names = ["", "", ""]

    filename = os.path.join(path_assets, "assets/materials/default.materials")
    if os.path.exists(filename) == False:
        return texture_names

    file = open(filename, 'r')

    # Parse material (very primitive ...)
    material_scope = False
    material_found = False
    for line in file.readlines():
        tokens = line.split()
        num = len(tokens)
        if num <= 0:
            continue
        token0 = tokens[0].strip()
        if material_scope == True:
            if token0 == "}":
                if material_found == True:
                    break
                material_scope = False
                continue
            if num < 3:
                continue
            token1 = tokens[1].strip()
            token2 = tokens[2].strip()

            if material_found == True:
                if (token0 == "DiffuseMap" or token0 == "SpecularMap" or token0 == "NormalMap") and token1 == "=":
                    if token2.endswith(','):
                        token2 = token2[0:-1]
                    if token2.startswith('"') and token2.endswith('"'):
                        token2 = token2[1:-1]
                    if token0 == "DiffuseMap":
                        texture_names[0] = token2
                    elif token0 == "NormalMap":
                        texture_names[1] = token2
                    elif token0 == "SpecularMap":
                        texture_names[2] = token2
            else:
                if token0 == "Name" and token1 == "=":
                    if token2.endswith(','):
                        token2 = token2[0:-1]
                    if token2.startswith('"') and token2.endswith('"'):
                        token2 = token2[1:-1]
                    if token2 == find_material:
                        material_found = True
        else:
            if token0 == "material{":
                material_scope = True
            elif num > 1 and (token0 == "material" and tokens[1] == "{"):
                material_scope = True
    file.close()

    for i in range(len(texture_names)):
        texture_name = texture_names[i]
        if len(texture_name) <= 0:
            continue

        # Check if source exists (.tga) or whatever was given in material file
        source_name = os.path.join(path_assets, texture_name)
        if os.path.exists(source_name):
            texture_name[i] = source_name
            continue

        # Check for .d3d9_texture instead
        base_name, ext = os.path.splitext(source_name)
        dxt_name = base_name + ".d3d9_texture"

        if os.path.exists(dxt_name):
            texture_names[i] = dxt_name
            continue

        # Didn't find file, just remove it
        texture_names[i] = ""

    return texture_names


def load_mesh(filename, context):
    # Dig out file base name and extension
    name, ext = os.path.splitext(os.path.basename(filename))

    path_assets = build_assets_path(filename)

    print("Opening file: " + filename)
    file = open(filename, 'rb')
    try:
        magic = struct.unpack("<4s", file.read(4))[0]
    except:
        print("Error parsing file header!")
        file.close()
        return

    # Figure out if it's a valid mesh file
    if magic != b'MESH':
        print("Not a valid mesh model!")
        file.close()
        return

    mesh_unknown = read_int(file)
    num_vertices = read_int(file)

    # This will store all vertices and index info
    vertex_set = []
    indices = []
    materials = []

    for i in range(15):
        # Assumption that this is what it actually means ...
        data_type = read_int(file)
        num_comp = read_int(file)
        byte_width = read_int(file)

        print( "Vertex Attrib %d" % i )
        print( "\tData Type: %d" % data_type )
        print( "\tNum Components: %d" % num_comp )
        print( "\tByte Width: %d" % byte_width )

        # Skip empty
        if data_type == 0 and num_comp == 0 and byte_width == 0:
            # Add empty set
            vertex_set.append([])
            continue

        # Report unknown data types (?)
        type_size = 0
        if data_type == 0:
            type_size = 1   # byte ?
        elif data_type == 2:
            type_size = 4   # int32 ?
        elif data_type == 3:
            type_size = 4   # float
        else:
            print("Valid mesh, but not supported vertex data: data_type = %d" % data_type)
            file.close()
            return

        # byte_width should be type_size*num_comp
        if byte_width != (num_comp*type_size):
            print("Valid mesh, but not supported vertex data: data_size = %d" % data_size)
            file.close()
            return

        vertex_data = []
        for j in range(num_vertices):
            data = []
            if num_comp == 2:
                if data_type == 0:
                    data = read_byte2(file)
                elif data_type == 2:
                    data = read_int2(file)
                elif data_type == 3:
                    data = read_float2(file)
            elif num_comp == 3:
                if data_type == 0:
                    data = read_byte3(file)
                elif data_type == 2:
                    data = read_int3(file)
                elif data_type == 3:
                    data = read_float3(file)
            elif num_comp == 4:
                if data_type == 0:
                    data = read_byte4(file)
                elif data_type == 2:
                    data = read_int4(file)
                elif data_type == 3:
                    data = read_float4(file)
            vertex_data.append(data)

        vertex_set.append(vertex_data)

    num_indices = read_int(file)
    for i in range(num_indices):
        index = read_int(file)
        indices.append(index)

    num_materials = read_int(file)
    for i in range(num_materials):
        name_length = read_int(file)
        material_name = read_string(file, name_length)
        print("Material Name: " + material_name)

        material_unknown = read_int(file)
        #if pre_material != 2:
        #    print("Valid mesh, but not supported material data: pre_material = " + str(pre_material))
        #    file.close()
        #    return

        index_offset = read_int(file)
        num_primitives = read_int(file)

        material = _mesh_material()
        material.name = material_name
        material.index_offset = index_offset
        material.num_primitives = num_primitives

        materials.append(material)

    # Origin x,y,z (?)
    origin = read_float3(file)

    # ?
    read_float(file)

    # Bounds min x,y,z (?)
    bounds_min = read_float3( file )

    # Bounds max x,y,z (?)
    bounds_max = read_float3( file )

    file.close()

    # Build the blender object
    build_objects(name, vertex_set, indices, materials, path_assets)


def add_material_texture( bl_material, type, filename ):
    # Create texture
    bl_texture = bpy.data.textures.new(name=type, type="IMAGE")

    # Try to load image
    image = load_image(filename)
    if image:
        bl_texture.image = image

    # Create texture slot and link to 'uvset1'
    mtex = bl_material.texture_slots.add()
    mtex.texture = bl_texture
    mtex.texture_coords = "UV"

    # Ensure diffuse/specular/normal uses proper types
    mtex.use_map_color_diffuse = (type == "diffuse")
    mtex.use_map_color_spec = (type == "specular")
    mtex.use_map_normal = (type == "normal")

    return image


def create_material( path_assets, material ):
    # Parse material file for the actual texture images
    texture_names = parse_material_file(path_assets, material.name)

    # Create blender material
    bl_material = bpy.data.materials.new(material.name)

    # Create textures and texture slots on material
    image_diffuse = add_material_texture(bl_material, "diffuse", texture_names[0])
    image_normal = add_material_texture(bl_material, "normal", texture_names[1])
    image_specular = add_material_texture(bl_material, "specular", texture_names[2])

    # Assign material
    material.bl_mat = bl_material

    # Assign preview display image for material
    material.bl_image = image_diffuse


def build_objects(name, vertex_set, indices, materials, path_assets):
    print("Building Blender data")
    print("Vertex sets: %d" % len(vertex_set))

    vertices = vertex_set[0]
    normals = vertex_set[1]
    tangents = vertex_set[2]
    bitangents = vertex_set[3]
    byte_colors = vertex_set[4]
    texcoords = vertex_set[5]

    num_vertices = len(vertices)
    num_indices = len(indices)
    num_primitives = int(num_indices / 3)
    num_materials = len(materials)

    print("Num vertices: " + str(num_vertices))
    print("Num indices: " + str(num_indices))
    print("Num primitives: " + str(num_primitives))
    print("Num materials: " + str(num_materials))

    # Convert color attribute from [0, 255] to [0, 1] range
    float_colors = []
    if len(byte_colors) > 0:
        for i in range(num_vertices):
            r = byte_colors[i][0] / 255.0
            g = byte_colors[i][1] / 255.0
            b = byte_colors[i][2] / 255.0
            a = byte_colors[i][3] / 255.0
            float_colors.append([r, g, b, a])
        print(float_colors)

    # Before adding any meshes or armatures go into Object mode.
    if bpy.ops.object.mode_set.poll():
        bpy.ops.object.mode_set(mode='OBJECT')

    # Create mesh
    me = bpy.data.meshes.new(name)

    # Add vertices
    me.vertices.add(num_vertices)
    for i in range(num_vertices):
        # Flip y-z component so it looks more natural in blender
        me.vertices[i].co[0] = vertices[i][0]
        me.vertices[i].co[1] = vertices[i][2]
        me.vertices[i].co[2] = vertices[i][1]

    # Add faces
    me.faces.add(num_primitives)
    for fi in range(num_primitives):
        idx = fi * 3
        for i in range(3):
            # Flip indices because of vertex flip
            me.faces[fi].vertices_raw[2-i] = indices[idx+i]

    # Add vertex colors
    if len(float_colors) > 0:
        clrmap = me.vertex_colors.new("clrset1")
        for fi in range(num_primitives):
            clrf = clrmap.data[fi]
            clrf.color1 = float_colors[indices[(fi*3)+0]][0:3]
            clrf.color2 = float_colors[indices[(fi*3)+0]][0:3]
            clrf.color3 = float_colors[indices[(fi*3)+0]][0:3]

    # Add uv map
    uvmap = me.uv_textures.new("uvset1")
    for fi in range(num_primitives):
        uvf = uvmap.data[fi]
        uvf.uv1 = texcoords[indices[(fi*3)+0]]
        uvf.uv2 = texcoords[indices[(fi*3)+1]]
        uvf.uv3 = texcoords[indices[(fi*3)+2]]

        # Flip uv coordinates
        uvf.uv1.y = 1.0 - uvf.uv1.y
        uvf.uv2.y = 1.0 - uvf.uv2.y
        uvf.uv3.y = 1.0 - uvf.uv3.y

    # Create object and link with scene
    ob = bpy.data.objects.new(name, me)
    bpy.context.scene.objects.link(ob)

    # Add/create all our materials
    for material in materials:
        create_material(path_assets, material)

    # Assign materials to mesh
    for mi in range(num_materials):
        material = materials[mi]
        me.materials.append(material.bl_mat)

        face_offset = int(material.index_offset / 3)
        for fi in range(material.num_primitives):
            me.faces[face_offset+fi].material_index = mi

            uvf = uvmap.data[face_offset+fi]
            uvf.image = material.bl_image

    # Update mesh and scene
    me.update()
    bpy.context.scene.update()


class IMPORT_OT_mesh(bpy.types.Operator, ImportHelper):
    # Import Mesh Operator.
    bl_idname = "import_scene.mesh"
    bl_label = "Import Mesh"
    bl_description = "Import a Legend of Grimrock mesh"
    bl_options = { 'REGISTER', 'UNDO' }

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

    def execute(self, context):
        load_mesh(self.filepath, context)
        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(IMPORT_OT_mesh.bl_idname, text="Legend of Grimrock Mesh (.mesh)")


def register():
    bpy.utils.register_module(__name__)
    bpy.types.INFO_MT_file_import.append(menu_func)


def unregister():
    bpy.utils.unregister_module(__name__)
    bpy.types.INFO_MT_file_import.remove(menu_func)


if __name__ == "__main__":
    register()