# 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
from mathutils import Vector
from mathutils import Quaternion
from shutil import copyfile
from . import kex_utils
from . import turok_export_anim
from bpy_extras.wm_utils.progress_report import (
    ProgressReport,
    ProgressReportSubstep,
)

# -----------------------------------------------------------------------------
#
def save(context,
         filepath,
         *,
         use_selection=True,
         global_matrix=None,
         path_mode='AUTO',
         exportMaterials=False,
         copyTextures=False,
         copyTextureSets=False,
         exportAnim=False,
         animNLATracks=False,
         useModelsDir=False
         ):

    obj = context.object
    #check obj parents for an armature object and use that instead if found
    while True:
        if obj.parent is None or obj.parent.type != "ARMATURE":
            break
        obj = obj.parent
    #end while
    
    validateResults = kex_utils.T1ValidateModel(context, obj, False)
    if len(validateResults["errors"]) > 0:
        return {'CANCELLED'}
    #end if

    nodeCount = 1
    if obj.type == "ARMATURE":
        nodeCount = len(obj.data.bones)
    #end if
    armatureChildren = obj.children
    with ProgressReport(context.window_manager) as progress:
        stepCount = 1
        if exportMaterials: stepCount += 1
        if copyTextures: stepCount += 1
        progress.enter_substeps(stepCount)
        print("Exporting Turok Model %r ..." % filepath)

        fileDir, fileName = os.path.split(filepath)
        fileNameNoExt = os.path.splitext(fileName)[0]
        
        matTexAnimDir = fileDir
        fileSubDir = ""
        matSubDir = "materials"
        if useModelsDir:
            matTexAnimDir, fileSubDir = kex_utils.GetPathFromDirName(fileDir, "models")
            #if found models directory
            if matTexAnimDir != fileDir and fileSubDir != "":
                matSubDir = "materials/%s" % kex_utils.GetLocalKexPathFromDir(fileSubDir)
            #end if
        #end if

        scene = bpy.context.scene
        dictKMats = {} #relative .kmat filepath, KMaterial
        validMorphActor = (obj.type == "ARMATURE")
        
        #Save the model.bin file
        with open(filepath, 'wb') as data:
            data.seek(0)
            kex_utils.write32(data, 1, False) #version (1)
            kex_utils.writeVector(data, obj.t1.boundingBoxMin) #bounding box min
            kex_utils.writeVector(data, obj.t1.boundingBoxMax) #bounding box max
            kex_utils.write32(data, nodeCount, False)
            if obj.type == "ARMATURE":
                #get list of armature bones and sort them by Turok boneIndex
                bones = []
                for bone in obj.data.bones:
                    bones.append(bone)
                bones.sort(key=lambda x: x.t1.boneIndex)
                
                if len(bones) > 1:
                    validMorphActor = False

                #get max variants this model uses on all bones
                variantLength = 1
                for bone in bones:
                    if len(bone.t1.variants) > variantLength:
                        variantLength = len(bone.t1.variants)

                for i in range(len(bones)):
                    boneIndex = i
                    bone = bones[boneIndex]
                    boneChildren = bone.children
                    
                    #get all objects(meshes) that belong to this bone (and sort them by variantIndex)
                    boneObjects = []
                    for i2 in range(len(armatureChildren)):
                        if bone.name == armatureChildren[i2].parent_bone and armatureChildren[i2].type == 'MESH':
                            boneObjects.append(armatureChildren[i2])
                    boneObjects.sort(key=lambda x: x.t1.variantIndex)
                    
                    kex_utils.write32(data, len(boneChildren), False)
                    kex_utils.write32(data, variantLength, False)
                    kex_utils.write32(data, kex_utils.max(1, len(boneObjects)), False)
                    
                    boneChildrenIndexs = []
                    for i2 in range(len(boneChildren)):
                        boneChildrenIndexs.append(boneChildren[i2].t1.boneIndex)
                    boneChildrenIndexs.sort()
                    
                    for boneChildIndex in boneChildrenIndexs:
                        kex_utils.write16(data, boneChildIndex, False)
                        
                    for i2 in range(variantLength):
                        variantIndex = 0
                        if i2 < len(bone.t1.variants):
                            variantIndex = bone.t1.variants[i2].variantIndex
                        kex_utils.write16(data, variantIndex, False)
                    
                    #if has no objects for this bone then write one with no surfaces
                    if len(boneObjects) == 0:
                        kex_utils.write32(data, 0, False) #surface count
                        continue

                    meshVertCount = 0
                    meshTriCount = 0
                    for i2 in range(len(boneObjects)):
                        boneObj = boneObjects[i2]
                        mVertCount = 0
                        mTriCount = 0
                        
                        #go through materials and check which materials are actually used in the model, return that new material list
                        materials = get_obj_materials(boneObj)
                        for mat in materials:
                            add_obj_material_to_kmatdict(fileNameNoExt, mat, dictKMats, matSubDir)
                        
                        kex_utils.write32(data, len(materials), False) #surface count
                        if len(materials) != 0:
                            mVertCount, mTriCount = write_mesh_surfaces(boneObj, data, fileNameNoExt, materials)
                        #end if
                        
                        #Check if is valid morph actor
                        if validMorphActor:
                            if i2 == 0:
                                meshVertCount = mVertCount
                                meshTriCount = mTriCount
                            elif meshVertCount != mVertCount or meshTriCount != mTriCount:
                                validMorphActor = False
                            #end if
                        #end if
                    #end for
                #end for
            else: #is a mesh
                kex_utils.write32(data, 0, False) #childrenLength
                kex_utils.write32(data, 1, False) #variantsLength
                kex_utils.write32(data, 1, False) #objectsLength
                kex_utils.write16(data, 0, False) #variant
                #write object
                materials = get_obj_materials(obj)
                for mat in materials:
                    add_obj_material_to_kmatdict(fileNameNoExt, mat, dictKMats, matSubDir)
                
                kex_utils.write32(data, len(materials), False) #surface count
                if len(materials) > 0:
                    write_mesh_surfaces(obj, data, fileNameNoExt, materials)
            #end if
        #end with
        
        extraInfo = ""
        if validMorphActor:
            extraInfo = " (Is a valid _morph animation actor)"
            
        progress.step("Model Successfully Exported!%s" % (extraInfo))

        if exportMaterials:
            #read existing files and put all those materials that aren't already in the dictionary in it
            for kmatPath, kMaterials in dictKMats.items(): 
                newKmatPath = kex_utils.find_t1_file_path(kmatPath + ".kmat")
                if newKmatPath is None:
                    continue
                newkMaterials = kex_utils.T1ReadKMat(newKmatPath)
                if newkMaterials is None:
                    continue
                for newKmatName, newKmat in newkMaterials.items(): 
                    hasMat = False
                    for kmat in kMaterials:
                        if newKmatName == kmat.name:
                            hasMat = True
                            break
                        #end if
                    #end for
                    if not hasMat:
                        dictKMats[kmatPath].append(newKmat)
                    #end if
                #end for
            #end for

            #write kmat files to directory relative to export filepath
            matCount = 0
            matFileCount = 0
            for kmatPath, kMaterials in dictKMats.items(): 
                matFileCount += 1
                matCount += len(kMaterials)
                kmatFilepath = os.path.normpath(os.path.join(matTexAnimDir, kmatPath + ".kmat"))
                kmatDir = os.path.split(kmatFilepath)[0]
                if not os.path.isdir(kmatDir):
                    os.makedirs(kmatDir)
                #end if
                with open(kmatFilepath, 'w+t', encoding='utf-8') as file:
                    for kmat in kMaterials:
                        file.write("material %s {\n" % (kmat.name))
                        file.write('\tshader "%s"\n' % (kmat.shaderPath))
                        file.write("\tsort %s\n" % (kmat.sort))
                        if kmat.cull != "none":
                            file.write("\tcull %s\n" % (kmat.cull))
                        if kmat.blend:
                            file.write("\tblend\n")
                        if kmat.fullBright:
                            file.write("\tfullbright\n")
                        if kmat.alphaTest:
                            file.write("\talphatest\n")
                            file.write("\talphafunc %s\n" % (kmat.alphaFunc))
                            file.write("\talphamask %f\n" % (kmat.alphaMask))
                        if kmat.noDraw:
                            file.write("\tnodraw\n")
                        if kmat.depthTest:
                            file.write("\tdepthtest\n")
                        if kmat.noDepthMask:
                            file.write("\tnodepthmask\n")
                        file.write('\tcolor_diffuse "%f %f %f %f"\n' % (kmat.diffuseColor[0], kmat.diffuseColor[1], kmat.diffuseColor[2], kmat.diffuseColor[3]))
                        for param in kmat.params:
                            file.write('\tparam "%s" %i "%s"\n' % (param.name, param.valueType, param.value))
                        file.write('\n')
                        kmat.textureInfo.sort()
                        for texInfoIndex, texInfo in enumerate(kmat.textureInfo):
                            if texInfoIndex > 3: #can only have 4 textures per material
                                break
                            file.write('\tsampler "%s" %i "%s" {\n' % (texInfo.type, texInfo.index, texInfo.path))
                            file.write('\t\tfilter\t\t%s\n' % (texInfo.filter))
                            file.write('\t\twrap\t\t%s\n' % (texInfo.wrap))
                            file.write('\t}\n')
                        #end for
                        file.write('}\n\n')
                    #end for
                #end with
            #end for
            progress.step("%i .kmat file(s) Exported (Materials: %i)" % (matFileCount, matCount))
        #end if
        
        if copyTextures:
            #get all used materials
            usedMaterials = []
            if obj.type == "ARMATURE":
                for childObj in armatureChildren:
                    mats = get_obj_materials(childObj)
                    for mat in mats:
                        if mat not in usedMaterials:
                            usedMaterials.append(mat)
                        #end if
                    #end for
                #end for
            else: #staticmesh
                usedMaterials = get_obj_materials(obj)
            #end if
            
            #get all used texture images
            images = {} #relative tex path, texture image
            texSetsToCopy = [] #array of string (tex set key)
            for mat in usedMaterials:
                if copyTextureSets:
                    #get first material texture set index
                    firstMat = kex_utils.get_material_texset_index(mat, 0, False)

                    #go through nodes of first mat and add them to texSetsToCopy
                    for node in firstMat.node_tree.nodes:
                        if node.bl_idname == 'ShaderNodeTexImage':
                            #save image textures using the filepath if not using a texture set
                            if node.image is not None:
                                localTexPath = kex_utils.texture_filepath_to_local_path(node.image.filepath)
                                if localTexPath not in kex_utils.turok1BuiltinTextures:
                                    texSet = kex_utils.get_textureset(localTexPath)
                                    if texSet is not None and localTexPath not in texSetsToCopy:
                                        texSetsToCopy.append(localTexPath)
                                    #end if
                                #end if
                            #end if
                        #end if
                    #end for
                #end if
                
                for node in mat.node_tree.nodes:
                    if node.bl_idname == 'ShaderNodeTexImage':
                        #save image textures using the filepath if not using a texture set
                        if node.image is not None:
                            localTexPath = kex_utils.texture_filepath_to_local_path(node.image.filepath)
                            if localTexPath not in kex_utils.turok1BuiltinTextures:
                                images[localTexPath] = node.image
                            #end if
                        #end if
                    #end if
                #end for
            #end for

            #save used images
            for localTexPath, image in images.items():
                texFilepath = os.path.join(matTexAnimDir, localTexPath)
                texDir = os.path.split(texFilepath)[0]
                if not os.path.isdir(texDir):
                    os.makedirs(texDir)
                #end if
                image.save_render(texFilepath)
            #end for
            
            #copy texture set local path texture files
            texsCopied = 0
            for texSetKey in texSetsToCopy:
                texSet = kex_utils.get_textureset(texSetKey)
                for texLocalPath in texSet.textures:
                    #don't copy textures that were already saved with images
                    if texLocalPath in images:
                        continue
                    texSourcePath = kex_utils.find_t1_file_path(texLocalPath)
                    if texSourcePath is None:
                        continue
                        
                    texDestPath = os.path.join(matTexAnimDir, texLocalPath)
                    try:
                        os.makedirs(os.path.dirname(texDestPath), exist_ok=True)
                        copyfile(texSourcePath, texDestPath)
                        texsCopied += 1
                    except:
                        pass
                    #end try except
                #end if
            #end for
            
            progress.step("%i Textures Exported" % (len(images) + texsCopied))
        #end if
        
        if exportAnim and obj.type == "ARMATURE":
            animDir = fileDir
            if useModelsDir and matTexAnimDir != fileDir:
                animSubDir = fileSubDir
                if animSubDir == "":
                    animSubDir = "anims"
                else:
                    animSubDir = os.path.join("anims", animSubDir)
                #end if
                animDir = os.path.join(matTexAnimDir, animSubDir)
                if not os.path.isdir(animDir):
                    os.makedirs(animDir)
                #end if
            #end if
            
            turok_export_anim.export_anim(os.path.join(animDir, fileNameNoExt + "_anim.bin"), context, animNLATracks)
        #end if
        
    #end with
        
    return {'FINISHED'}
#end func
# -----------------------------------------------------------------------------
# Goes through materials and checks which materials are actually used in the obj mesh
# and returns that new material list.
def get_obj_materials(obj):
    materialIndexsUsed = []
    matCount = len(obj.data.materials)
    for poly in obj.data.polygons:
        if poly.material_index not in materialIndexsUsed and poly.material_index < matCount:
            materialIndexsUsed.append(poly.material_index)
            if len(materialIndexsUsed) == matCount: #used all materials
                break
            #end if
        #end if
    #end for
    
    return [material for matIndex, material in enumerate(obj.data.materials) if material is not None and matIndex in materialIndexsUsed]
#end func

# -----------------------------------------------------------------------------
# 
def add_obj_material_to_kmatdict(fileNameNoExt, mat, dict, matSubDir):
    if mat.t1.isMatPathCustom:
        key = mat.t1.matPath
    else:
        key = "%s/mat_%s" % (matSubDir, fileNameNoExt)
    #end if
    if key not in dict:
        dict[key] = []
    hasKmat = False
    matName = kex_utils.get_real_material_name(mat.name)[0]
    for kmat in dict[key]:
        if matName == kmat.name:
            hasKmat = True
            break
        #end if
    #end for
    if not hasKmat:
        dict[key].append(kex_utils.material_to_kmat(mat))
    #end if
#end func

# -----------------------------------------------------------------------------
# Returns total amount of verts, tris used by all surfaces
def write_mesh_surfaces(obj, data, fileNameNoExt, materials):
    totalVertCount = 0
    totalTriCount = 0
    mesh = obj.data
    
    for material in materials:
        matIndex = kex_utils.find_material_index(mesh, material)
        kMatPath = ""
        matName = kex_utils.get_real_material_name(material.name)[0]
        if material.t1.isMatPathCustom:
            kMatPath = "%s/%s" % (material.t1.matPath, matName)
        else:
            kMatPath = "materials/mat_%s/%s" % (fileNameNoExt, matName)
        kex_utils.writeString(data, kMatPath) #surface name (e.g. materials/mat_ammo_bullet_box01/surf_113_0000A)
        
        #get/add the triangle indices, and verts that this material uses
        tris = []
        verts = []
        uvDict = {} #key is tuple, value is vert index for local verts[]
        uv_layer = mesh.uv_layers.active.data
        tupUV = None
        meshVert = None
        for poly in mesh.polygons:
            if poly.material_index == matIndex:
                tri = []
                for loopIndex in range(poly.loop_start, poly.loop_start + poly.loop_total):
                    meshVert = mesh.vertices[mesh.loops[loopIndex].vertex_index]
                    tupUV = tuple([meshVert.co[0], meshVert.co[1], meshVert.co[2], uv_layer[loopIndex].uv[0], uv_layer[loopIndex].uv[1], meshVert.normal[0], meshVert.normal[1], meshVert.normal[2]])
                    if tupUV in uvDict:
                        tri.append(uvDict[tupUV])
                    else:
                        tri.append(len(verts))
                        uvDict[tupUV] = len(verts)
                        verts.append((meshVert.co, uv_layer[loopIndex].uv, meshVert.normal))
                    #end if
                #end for
                tris.append(tri)
            #end if
        #end if
        
        totalVertCount += len(tris)
        totalTriCount += len(verts)

        #each material(surface) has a list of there own triangles and there own vertices
        kex_utils.write32(data, len(tris), False) #tri count
        for tri in tris:
            for indice in tri:
                kex_utils.write16(data, indice, False)
            #end for
        #end for
        
        kex_utils.write32(data, len(verts), False) #vertice count
        for vert in verts:
            kex_utils.writeVector(data, Vector(vert[0]))
            kex_utils.writeVector2(data, vert[1])
            kex_utils.writeVector(data, Vector(vert[2]))
        #end for
    #end for
    return totalVertCount, totalTriCount
#end func
