# 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.
#

import os
import bpy
import bmesh
import struct

from mathutils import Vector
from mathutils import Quaternion
from bpy_extras.image_utils import load_image
from . import kex_utils
from . import turok_material_panel
from bpy_extras.wm_utils.progress_report import (
    ProgressReport,
    ProgressReportSubstep,
)

# -----------------------------------------------------------------------------
# Returns dictionary of model structure. Returns None if there was an error.
def load_model_bin(filepath, global_matrix):
    model = {}
    with open(filepath, 'rb') as data:
        data.seek(0)
        model["version"] = kex_utils.read32(data, False) #version (1)
        if model["version"] != 1:
            kex_utils.show_error("[Import Model] %s has incorrect version number of %d. Version must be 1." % (filepath, model["version"]))
            return None
        model["bbMin"] = kex_utils.readVector(data) #bounding box min
        model["bbMax"] = kex_utils.readVector(data) #bounding box max
        model["totalTriCount"] = 0
        model["totalVertCount"] = 0
        model["totalMeshCount"] = 0
        model["nodes"] = []
        model["unique_kmats"] = {}
        
        numNodes = kex_utils.read32(data, False)
        if numNodes <= 0:
            kex_utils.show_error("[Import Model] %s has no nodes. There should always be at least 1 node." % (filepath))
            return None
        if numNodes > 65535:
            kex_utils.show_error("[Import Model] %s has %d nodes. 65535 is the node limit." % (filepath, numNodes))
            return None
            
        for i in range(numNodes):
            numChildren = kex_utils.read32(data, False)
            if numChildren > 65535:
                kex_utils.show_error("[Import Model] %s has %d children nodes. There should never be more than 65535 children." % (filepath, numChildren))
                return None
                
            numVariant = kex_utils.read32(data, False)
            if numVariant <= 0:
                kex_utils.show_error("[Import Model] %s Node %d has no variants. There should always be at least 1 variant in a node." % (filepath, i))
                return None
            elif numVariant > 255:
                kex_utils.show_error("[Import Model] %s Node %d has %d variants. The limit is 255 variants." % (filepath, i, numVariant))
                return None
                
            numObjects = kex_utils.read32(data, False)
            if numObjects <= 0:
                kex_utils.show_error("[Import Model] %s Node %d has no objects. There should always be at least 1 object in a node." % (filepath, i))
                return None
            elif numObjects > 255:
                kex_utils.show_error("[Import Model] %s Node %d has %d objects. The limit is 255 objects." % (filepath, i, numObjects))
                return None
            
            node = { "children": [0] * numChildren, "variations": [0] * numVariant, "objects": [{}] * numObjects }
            model["nodes"].append(node)
            
            for j in range(numChildren):
                node['children'][j] = kex_utils.read16(data, False)
            #end if
                
            for j in range(numVariant):
                node['variations'][j] = kex_utils.read16(data, False)
            #end if
                
            for j in range(numObjects):
                numSurfaces = kex_utils.read32(data, False)
                node['objects'][j] = {}
                object = node['objects'][j]
                object["surfaces"] = []
                object["indices"] = []
                object["vertices"] = []
                object["coords"] = []
                object["normals"] = []
                object["bObj"] = None
                if numSurfaces == 0:
                    continue
                    
                nTotalTris = 0
                nCurTris = 0
                for k in range(numSurfaces):
                    surface = { "material": "", "tris": [] }
                    object['surfaces'].append(surface)
                    
                    surface["material"] = kex_utils.readString(data)
                        
                    #Triangles
                    numTris = kex_utils.read32(data, False)
                    model["totalTriCount"] += numTris
                    if numTris == 0:
                        continue
                    for t in range(numTris):
                        v0 = kex_utils.read16(data, False) + nTotalTris
                        v1 = kex_utils.read16(data, False) + nTotalTris
                        v2 = kex_utils.read16(data, False) + nTotalTris
                        
                        tri = (v0, v1, v2)
                        surface["tris"].append(tri)
                        object["indices"].append(tri)
                        nCurTris = kex_utils.max(kex_utils.max(kex_utils.max(nCurTris, v0), v1), v2);
                    #end for
                    
                    #Vertices
                    nTotalTris = nCurTris + 1;
                    numVert = kex_utils.read32(data, False)
                    model["totalVertCount"] += numVert
                    if numVert == 0:
                        continue
                    for v in range(numVert):
                        object["vertices"].append(global_matrix @ kex_utils.readVector(data))
                        object["coords"].append((kex_utils.readFloat(data), kex_utils.readFloat(data)))
                        object["normals"].append(kex_utils.readVector(data))
                    #end for
                    
                    # keep track of what kmat files are being used and what materials in those files are being used. 
                    kmatRelPath, kMaterialName = os.path.split(surface["material"])
                    if kmatRelPath not in model["unique_kmats"]:
                        model["unique_kmats"][kmatRelPath] = []
                    if kMaterialName not in model["unique_kmats"][kmatRelPath]:
                        model["unique_kmats"][kmatRelPath].append(kMaterialName)
                        
                #end for
            #end for
        #end for
    #end with
    
    return model
#end func

# -----------------------------------------------------------------------------
# returns rootObj back
def create_mesh_object(rootObj, object, nodeID, objectID, fileNameNoExt, useMapModelType=0):
    mesh = bpy.data.meshes.new('mesh')
    #create mesh from tris and verts
    if len(object["indices"]) > 0 and len(object["vertices"]) > 0:
        mesh.from_pydata(object["vertices"], [], object["indices"])
        mesh.validate(clean_customdata=False)
        mesh.update(calc_edges=True)

        #set custom normals
        mesh.use_auto_smooth = True
        mesh.normals_split_custom_set_from_vertices(object["normals"])
        mesh.calc_normals_split()
        for vert in mesh.vertices:
            vert.normal = object["normals"][vert.index]
        for loop in mesh.loops:
            loop.normal = object["normals"][loop.vertex_index]

        # create uv textures
        uv_layer = mesh.uv_layers.new()
        for loop in mesh.loops:
            uv_layer.data[loop.index].uv = object["coords"][loop.vertex_index]
        #end for
    #end if

    # create mesh scene object
    meshObjName = 'mesh_%03d_%03d' % (nodeID, objectID) #mesh_nodeID_objectID
    turokObjectType = "MODEL"
    if useMapModelType > 0:
        turokObjectType = "NONE"
    #end if
    obj = kex_utils.create_turok_object(meshObjName, mesh, turokObjectType)
    bpy.context.view_layer.objects.active = obj
    object["bObj"] = obj

    if rootObj is None:
        rootObj = obj
        rootObj.name = fileNameNoExt
    else:
        obj.parent = rootObj #parent mesh to the rootObj
    #end if
        
    #parent armature bone to mesh
    if rootObj.type == "ARMATURE":
        obj.parent_type = 'BONE'
        obj.parent_bone = kex_utils.get_bone_name(nodeID)
        obj.t1.variantIndex = objectID
    #end if
    
    return rootObj
#end func

# -----------------------------------------------------------------------------
# Returns rootObj if successful else None
def load(operator, progress, context, filepath, global_matrix, hideVariants, useMapModelType=0, cachedModel=None, unique_mats_override=None):
    rootObj = None
    
    if progress is not None and useMapModelType == 0:
        progress.enter_substeps(2)
        print("Importing Turok Model %r ..." % filepath)
    #end if

    fileName = os.path.split(filepath)[1]
    fileNameNoExt = os.path.splitext(fileName)[0]
    scene = bpy.context.scene
    scene.cursor.location = (0.0, 0.0, 0.0) #need to have the cursor position at center
    
    model = None
    if cachedModel is None:
        model = load_model_bin(filepath, global_matrix)
    else:
        model = cachedModel
    #end if
    
    if model is None:
        if useMapModelType > 0:
            return None, None, None
        #end if
        return None
    #end if
    
    unique_materials = {}
    if unique_mats_override is not None:
        unique_materials = unique_mats_override
    nodeCount = len(model["nodes"])
    
    if useMapModelType > 0:
        #Only the first node is created for staticmeshes
        if useMapModelType == 1 and len(model["nodes"]) > 1:
            model["nodes"] = [model["nodes"][0]]
            nodeCount = len(model["nodes"])
        #end if
        
        #Only the first object in the nodes are created
        for i in range(nodeCount):
            node = model["nodes"][i]
            
            if len(node['objects']) > 1:
                node['objects'] = [node['objects'][0]]
            #end with
        #end for
        
    #create an armature for the model if theres more than 1 node or has 1 node but more than 1 variant objects meshes
    elif nodeCount > 1 or len(model["nodes"][0]["objects"]) > 1:
        #Create armature and armature object for dynamic model
        armature = bpy.data.armatures.new(name='Armature')
        rootObj = kex_utils.create_turok_object(fileNameNoExt, armature, "MODEL")
        rootObj.data.display_type = 'STICK'
        rootObj.show_in_front = True            
        
        #Must make armature active and in edit mode to create edit bones
        bpy.context.view_layer.objects.active = rootObj
        kex_utils.mode_set('EDIT')
        
        #Create edit bones for each node
        for i in range(nodeCount):
            #Make a edit bone
            boneName = kex_utils.get_bone_name(i)
            bone = armature.edit_bones.new(boneName)
            bone.use_deform = False
            bone.head = (0.0, 0.0, 0.0)
            bone.tail = (0.0, 0.0000011, 0.0)
            bone.lock = True
        #end for
        
        #Exit Armature editing so the edit bone is created internally at this point
        kex_utils.mode_set('OBJECT')
    #end if
    
    meshesCreated = 0
    for i in range(nodeCount):
        node = model["nodes"][i]
        boneName = kex_utils.get_bone_name(i)
        objectCount = len(node['objects'])
        
        if rootObj is not None and rootObj.type == "ARMATURE":
            #set turok bone properties
            variantsCount = len(node["variations"])
            tBone = rootObj.data.bones[boneName].t1
            tBone.boneIndex = i
            for j in range(variantsCount):
                variant = tBone.variants.add()
                variant.name = "Variant %d" % (len(tBone.variants) - 1)
                variant.variantIndex = node["variations"][j]
            #end for
        #end if
        
        for j in range(objectCount):
            object = node['objects'][j]
            surfaceCount = len(object["surfaces"])
            if useMapModelType == 0 and (len(object["indices"]) == 0 or len(object["vertices"]) == 0):
                continue
            rootObj = create_mesh_object(rootObj, object, i, j, fileNameNoExt, useMapModelType)
            
            meshesCreated += 1
        #end for
    #end for
    
    #if there was no meshes created, then create at least 1 empty mesh on the root node
    #There is always at least 1 node and nodes always have at least 1 object.
    if meshesCreated == 0:
        rootObj = create_mesh_object(rootObj, model["nodes"][0]['objects'][0], 0, 0, fileNameNoExt, useMapModelType)
        meshesCreated += 1
    #end if
    
    model["totalMeshCount"] += meshesCreated

    #Make sure map version rootobj has correct objType
    if useMapModelType > 0 and rootObj is not None:
        if useMapModelType == 2:
            rootObj.t1.objType = "MAP ACTOR"
        else:
            rootObj.t1.objType = "MAP STATICMESH"
    #end if
    
    #set rootObj boundingBox
    if rootObj is not None and useMapModelType == 0:
        rootObj.t1.boundingBoxMin = model["bbMin"]
        rootObj.t1.boundingBoxMax = model["bbMax"]
    #end if

    if progress is not None and useMapModelType == 0:
        progress.step("Created Objects (verts:%i faces:%i meshes:%i bones:%i)" % (model["totalVertCount"], model["totalTriCount"], model["totalMeshCount"], nodeCount))
    #end if

    #parent objects for map model type
    mapModelObjs = []
    mapModelObjs.append(rootObj)
    if useMapModelType > 0 and nodeCount > 1 and rootObj is not None:
        for child in rootObj.children:
            mapModelObjs.append(child)
        #end for
        for i in range(nodeCount):
            if i >= len(mapModelObjs):
                break
            childList = model["nodes"][i]["children"]
            for j in range(len(childList)):
                childIdx = childList[j]
                if childIdx == i or childIdx >= len(mapModelObjs):
                    continue
                mapModelObjs[childIdx].parent = mapModelObjs[i]
            #end for
        #end for
        
    #parent the bones (see if can put this where the armature is created at the top in this function)
    elif nodeCount > 1 and rootObj is not None and rootObj.type == "ARMATURE":
        #Must make armature active and in edit mode to Edit edit bones
        bpy.context.view_layer.objects.active = rootObj
        kex_utils.mode_set('EDIT')
        
        for i in range(nodeCount):
            childList = model["nodes"][i]["children"]
            for j in range(len(childList)):
                childIdx = childList[j]
                if childIdx == i:
                    continue
                childBoneName = kex_utils.get_bone_name(childIdx)
                parentBoneName = kex_utils.get_bone_name(i)
                rootObj.data.edit_bones[childBoneName].parent = rootObj.data.edit_bones[parentBoneName]
            #end for
        #end for
        
        #Exit Armature editing
        kex_utils.mode_set('OBJECT')
    #end if
    
    #read kmat files and create blender materials and shader nodes and load images
    for kmatRelPath in model["unique_kmats"]:
        kmatPath = kex_utils.find_t1_file_path(kmatRelPath + ".kmat")
        if kmatPath is not None:
            kMaterials = kex_utils.T1ReadKMat(kmatPath)
            if kMaterials is None:
                print("Error reading kmat file: " + kmatPath)
            else:
                for kMaterialName in kMaterials:
                    if kMaterialName not in model["unique_kmats"][kmatRelPath]: #only include the materials that the model is using
                        continue
                        
                    kMaterialPath = kmatRelPath + "/" + kMaterialName
                    if kMaterialPath in unique_materials: #only add materials not already added to unique materials
                        continue
                    
                    kMatInfo = unique_materials[kMaterialPath] = kMaterials[kMaterialName] #assign unique kMatInfo 
                    matTexSetName = kex_utils.get_texset_material_name(kMaterialName)
                    kMatInfo.bMat = bpy.data.materials.get(kMaterialName)
                    if kMatInfo.bMat is None:
                        kMatInfo.bMat = bpy.data.materials.get(matTexSetName)
                        if kMatInfo.bMat is None:
                            kMatInfo.bMat = bpy.data.materials.new(kMaterialName)
                            kMatInfo.bMat.use_fake_user = True
                        #end if
                    #end if
                    kMatFilename = os.path.split(kmatRelPath)[1]
                    if kMatFilename == ("mat_" + fileNameNoExt):
                        kMatInfo.bMat.t1.isMatPathCustom = False
                    else:
                        kMatInfo.bMat.t1.isMatPathCustom = True
                    kMatInfo.bMat.t1.matPath = kmatRelPath
                    kex_utils.kmat_to_material(kMatInfo, kMatInfo.bMat)
                #end for
            #end if
        #end if
    #end for
    
    #assign materials to each surface
    for node in model["nodes"]:
        for object in node['objects']:
            if object["bObj"] is not None:
                for surface in object['surfaces']:
                    if surface["material"] in unique_materials:
                        bMat = unique_materials[surface["material"]].bMat
                        
                        #check if material is already in the mesh materials list
                        matSlotIndex = kex_utils.find_material_index(object["bObj"].data, bMat)
                        
                        #append material to mesh materials if not already in the mesh materials list
                        if matSlotIndex == -1:
                            matSlotIndex = len(object["bObj"].data.materials)
                            object["bObj"].data.materials.append(bMat)
                        #end if
                            
                        #assign all face materials to only the surface faces
                        for poly in object["bObj"].data.polygons:
                            polyVerts = (int(poly.vertices[0]), int(poly.vertices[1]), int(poly.vertices[2]))
                            if polyVerts in surface["tris"]:
                                poly.material_index = matSlotIndex
                        #end for
                    #end if
                #end for
            #end if
        #end for
    #end for
                                
    if hideVariants:
        for node in model["nodes"]:
            for i in range(1, len(node['objects'])):
                obj = node['objects'][i]["bObj"]
                if obj is not None:
                    obj.hide_set(state=True, view_layer=bpy.context.view_layer)
                #end if
            #end for
        #end for
    #end if
    
    if progress is not None and useMapModelType == 0:
        progress.leave_substeps("Model Successfully Imported!")
    #end if
    
    if useMapModelType > 0:
        return mapModelObjs, model, unique_materials
    else:
        return rootObj
#end func
