# 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 mathutils import Matrix, Euler

from enum import Enum
from bpy.props import FloatVectorProperty, StringProperty, PointerProperty
from . import kex_utils
from bpy_extras.wm_utils.progress_report import (
    ProgressReport,
    ProgressReportSubstep,
)

# -----------------------------------------------------------------------------
# Returns anim if successful else returns None
def load_anim(data, global_matrix):
    anim = {}
    anim["initialJoints"] = []
    anim["yawOffsets"] = []
    anim["actions"] = []
    
    anim["animID"] = kex_utils.read32(data, True)
    anim["numFrames"] = kex_utils.read32(data, False)
    if anim["numFrames"] < 2:
        kex_utils.show_error("[Import Anim] %s has %d frames. There should always be at least 2 frames." % (filepath, anim["numFrames"]))
        return None
    
    anim["numAnimsets"] = kex_utils.read32(data, False)
    if anim["numAnimsets"] > 65535:
        kex_utils.show_error("[Import Anim] %s has %d animSets. 65535 is the limit." % (filepath, anim["numAnimsets"]))
        return None
        
    anim["numNodes"] = kex_utils.read32(data, False)
    if anim["numNodes"] <= 0:
        kex_utils.show_error("[Import Anim] %s has no nodes. There should always be at least 1 node." % (filepath))
        return None
    if anim["numNodes"] > 65535:
        kex_utils.show_error("[Import Anim] %s has %d nodes. 65535 is the node limit." % (filepath, anim["numNodes"]))
        return None
            
    anim["numTranslations"] = kex_utils.read32(data, False)
    anim["numRotations"] = kex_utils.read32(data, False)
    anim["numActions"] = kex_utils.read32(data, False)
    anim["marker"] = kex_utils.read32(data, False)
    if anim["marker"] != 0:
        anim["blend"] = kex_utils.read16(data, False)
        anim["loopFrame"] = kex_utils.read16(data, False)
    else:
        anim["blend"] = 0
        anim["loopFrame"] = 0
        
    translationIndexes = []
    rotationIndexes = []
    
    for i in range(anim["numAnimsets"]):
        translationIndexes.append(kex_utils.read16(data, True))
        rotationIndexes.append(kex_utils.read16(data, True))
        
    anim["joints"] = [{}] * anim["numNodes"]
    anim["initialJoints"] = [{}] * anim["numNodes"]
    
    for i in range(anim["numNodes"]):
        vTranslation = kex_utils.readVector(data)
        vRotation = kex_utils.readQuat(data)
        anim["initialJoints"][i] = { "translation": Vector(global_matrix @ vTranslation), "rotation": vRotation}
        anim["joints"][i] = {}
        anim["joints"][i]["trans"] = [Vector((0.0, 0.0, 0.0))] * anim["numFrames"]
        anim["joints"][i]["rot"] = [Quaternion((1.0, 0.0, 0.0, 0.0))] * anim["numFrames"]
        
    for i in range(anim["numFrames"]):
        anim["yawOffsets"].append(kex_utils.readFloat(data)) #in rads
        
    translations = [[]] * anim["numTranslations"]
    rotations = [[]] * anim["numRotations"]
    
    for i in range(anim["numTranslations"]):
        translations[i] = [Vector((0.0, 0.0, 0.0))] * anim["numFrames"]
        
        for j in range(anim["numFrames"]):
            vTranslation = kex_utils.readVector(data)
            translations[i][j] = Vector(global_matrix @ vTranslation)
            
    for i in range(anim["numRotations"]):
        rotations[i] = [Quaternion((0.0, 0.0, 0.0, 0.0))] * anim["numFrames"]
        
        for j in range(anim["numFrames"]):
            rotations[i][j] = kex_utils.readQuat(data)
    
    for i in range(anim["numNodes"]):
        translationID = translationIndexes[i]
        rotationID = rotationIndexes[i]
        
        for j in range(anim["numFrames"]):
            pos = anim["initialJoints"][i]["translation"]
            rot = anim["initialJoints"][i]["rotation"]
            
            if translationID != -1:
                pos = translations[translationID][j]
                
            if rotationID != -1:
                rot = rotations[rotationID][j]
                
            anim["joints"][i]["trans"][j] = pos
            anim["joints"][i]["rot"][j] = rot
            
    for i in range(anim["numActions"]):
        anim["actions"].append(
        {
            "function": kex_utils.read32(data, True),
            "frame": kex_utils.read32(data, False),
            "arg1": kex_utils.readFloat(data),
            "arg2": kex_utils.readFloat(data),
            "arg3": kex_utils.readFloat(data),
            "arg4": kex_utils.readFloat(data)
        })
        
    return anim

# -----------------------------------------------------------------------------
# Returns list of anims if successful else returns None
def load_anim_bin(filepath, global_matrix, useMapVersion=False):
    with open(filepath, 'rb') as data:
        data.seek(0)
        version = kex_utils.read32(data, False) #version (1)
        if version != 1:
            kex_utils.show_error("[Import Anim] %s has incorrect version number of %d. Expected version 1." % (filepath, version))
            return None
        #end if
            
        animCount = kex_utils.read32(data, False)
        if animCount <= 0:
            kex_utils.show_error("[Import Anim] %s has no animations" % (filepath))
            return None
        #end if
        
        if useMapVersion:
            anim = load_anim(data, global_matrix)
            if anim is None:
                return None
            #end if
            
            return anim
        else:
            anims = [{}] * animCount
            for i in range(0, animCount):
                anims[i] = load_anim(data, global_matrix)
                if anims[i] is None:
                    return None
                #end if
            #end for
            return anims
        #end if
    #end with
#end func

# -----------------------------------------------------------------------------
# Returns the first anim data (None if error) and sets the objs position and rotation according to the anim.
# if animData is None then loads the anim file using filepath
def load_map_version(operator, context, objs, animData, filepath, global_matrix):
    from itertools import chain
    import math
    if animData is None:
        anim = load_anim_bin(filepath, global_matrix, useMapVersion=True)
    else:
        anim = animData
    #end if
    if anim is None:
        return None
    #end if

    for i in range(anim["numNodes"]):
        if i >= len(objs):
            break;
        if anim["numFrames"] > 0:
            pos = anim["joints"][i]["trans"][0]
            rot = anim["joints"][i]["rot"][0]
        else:
            pos = anim["initialJoints"][i]["translation"]
            rot = anim["initialJoints"][i]["rotation"]
        #end for
        objs[i].location = pos
        #poseBone.rotation_mode = 'QUATERNION'
        objs[i].rotation_quaternion = rot
    #end for
    
    return anim
#end func
# -----------------------------------------------------------------------------
# Returns true if successful
def load(operator, context, filepath, global_matrix, bakeYawOffsets, stashActions, anim60FPS):
    from itertools import chain
    import math
    
    fpsScale = (60.0 / kex_utils.turok1FPS) if anim60FPS else 1.0
    with ProgressReport(context.window_manager) as progress:
        progress.enter_substeps(1)
        print("Importing Turok Animation %r ..." % filepath)
        
        obj = context.object
        if obj is None:
            kex_utils.show_error("[Import Anim] No Object Selected")
            return False
            
        bpy.context.view_layer.objects.active = None
        scene = bpy.context.scene
        
        #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

        if obj.type != "ARMATURE" and obj.parent is not None:
            kex_utils.show_error("[Import Anim] Object (%s) must not be parented to anything as an Armature will be created and parented to this object." % (obj.name))
            return False
        
        anims = load_anim_bin(filepath, global_matrix)

        if anims is None:
            return False

        #get maximum number of nodes this animation uses
        usedNodesCount = 0
        for i in range(0, len(anims)):
            if anims[i]["numNodes"] > usedNodesCount:
                usedNodesCount = anims[i]["numNodes"]
                
        if usedNodesCount <= 0:
            kex_utils.show_error("[Import Anim] Animation (%s) uses no bones" % (filepath))
            return False
        
        #if not an armature (then has 1 node) and required anim nodes must be 1
        if obj.type != "ARMATURE":
            if usedNodesCount > 1:
                kex_utils.show_error("[Import Anim] Animation (%s) requires more than 1 bone. Object (%s) must be an Armature with at least %d bones to import this animation." % (filepath, obj.name, usedNodesCount))
                return False
            #create armature data and object and parent this object to that armature and create 1 bone
            armature = bpy.data.armatures.new(name='Armature')
            armatureObj = kex_utils.create_turok_object(obj.name, armature, "MODEL")
            armatureObj.data.display_type = 'STICK'
            armatureObj.show_in_front = True
            #Must make armature active and in edit mode to create edit bones
            bpy.context.view_layer.objects.active = armatureObj
            kex_utils.mode_set('EDIT')
            #Make a edit bone
            boneName = kex_utils.get_bone_name(0)
            bone = armatureObj.data.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
            #Exit Armature editing so the edit bone is created internally at this point
            kex_utils.mode_set('OBJECT')
            obj.parent = armatureObj
            obj.parent_type = 'BONE'
            obj.parent_bone = boneName
            armatureObj.t1.boundingBoxMin = obj.t1.boundingBoxMin
            armatureObj.t1.boundingBoxMax = obj.t1.boundingBoxMax
            obj = armatureObj
            tBone = armature.bones[boneName].t1
            tBone.boneIndex = 0
            variant = tBone.variants.add()
            variant.name = "Variant %d" % (len(tBone.variants) - 1)
            variant.variantIndex = 0
        #end if
        
        boneCount = len(obj.data.bones)
        if boneCount == 0:
            kex_utils.show_error("[Import Anim] Object (%s) has no bones. Animation (%s) requires %d bones" % (obj.name, filepath, usedNodesCount))
            return False
        
        if boneCount < usedNodesCount:
            kex_utils.show_error("[Import Anim] Animation (%s) requires %d bones. Object (%s) has only %d bones" % (filepath, usedNodesCount, obj.name, boneCount))
            return False
        
        #Must make armature active
        bpy.context.view_layer.objects.active = obj
        kex_utils.mode_set('POSE')
        animData = obj.animation_data_create()
        maxFrames = 0
        firstAction = None
        totalKeyFrameInserts = 0
        progress.enter_substeps(len(anims))
        for animIndex in range(0, len(anims)):
            anim = anims[animIndex]
            animName = kex_utils.get_animation_name(anim["animID"])

            animAction = bpy.data.actions.new(animName)
            animData.action = animAction
            tAction = animData.action.t1
            animAction.use_fake_user = True
            
            tAction.animID = anim["animID"]
            
            #Set key Frames
            for i in range(0, len(anim["actions"])):
                keyFrame = tAction.keyFrames.add()
                animKFAction = anim["actions"][i]
                keyFrame.frame = animKFAction["frame"]
                keyFrame.event = animKFAction["function"]
                keyFrame.arg1 = animKFAction["arg1"]
                keyFrame.arg2 = animKFAction["arg2"]
                keyFrame.arg3 = animKFAction["arg3"]
                keyFrame.arg4 = animKFAction["arg4"]
            #end
            
            #set marker, blend and loopframe
            tAction.marker = anim["marker"]
            tAction.blend = anim["blend"]
            tAction.loopFrame = anim["loopFrame"]
                
            if animIndex == 0:
                firstAction = animAction
                bpy.context.scene.frame_start = 0
                bpy.context.scene.frame_set(0)
            #end if
            
            hasYawOffset = False
            for frameIndex in range(0, anim["numFrames"]):
                if anim["yawOffsets"][frameIndex] != 0.0:
                    hasYawOffset = True
                    break

            if hasYawOffset:
                usedFrames = kex_utils.get_used_frames(anim["yawOffsets"])
                usedFramesLength = len(usedFrames)
                yawCurve = animAction.fcurves.new(data_path="t1.yawOffset")
                yawCurve.keyframe_points.add(count=usedFramesLength)
                for i in range(usedFramesLength):
                    frameIndex = usedFrames[i]
                    yawCurve.keyframe_points[i].co = (frameIndex * fpsScale, anim["yawOffsets"][frameIndex])
                    yawCurve.keyframe_points[i].interpolation = 'LINEAR'
                yawCurve.update()
                totalKeyFrameInserts += usedFramesLength
            #end if
            
            for i in range(0, boneCount):
                boneName = kex_utils.get_bone_name(i)
                if boneName not in obj.pose.bones:
                    kex_utils.show_error("[Import Anim] There is no pose bone named %s in the Armature" % (boneName))
                    return False

                #Get the new bone Translation and Rotation for all the frames
                boneTrans = []
                if i >= anim["numNodes"]:
                    boneTrans = [Vector((0.0, 0.0, 0.0))] * anim["numFrames"]
                else:
                    boneTrans = anim["joints"][i]["trans"]
                #end if
                boneRots = []
                if i >= anim["numNodes"]:
                    boneRots = [Quaternion((1.0, 0.0, 0.0, 0.0))] * anim["numFrames"]
                else:
                    if bakeYawOffsets and hasYawOffset and i == 0:
                        boneRots = anim["joints"][i]["rot"].copy()
                    else:
                        boneRots = anim["joints"][i]["rot"]
                    #end if
                #end if
                if bakeYawOffsets and hasYawOffset and i == 0:
                    for j in range(anim["numFrames"]):
                        boneRots[j] = anim["joints"][i]["rot"][j].rotate(Quaternion((0.0, 0.0, 1.0), anim["yawOffsets"][j]))

                #set initial poseBone position and rotation to first frame
                poseBone = obj.pose.bones[boneName]
                poseBone.location = boneTrans[0]
                poseBone.rotation_mode = 'QUATERNION'
                poseBone.rotation_quaternion = boneRots[0]

                #Create the curves for only the keyframes used
                usedLocFrames = [kex_utils.get_used_frames([v[j] for v in boneTrans]) for j in range(3)]
                usedRotFrames = [kex_utils.get_used_frames([q[j] for q in boneRots]) for j in range(4)]
                
                trans_curves = []
                for j in range(3):
                    trans_curves.append(animAction.fcurves.new(poseBone.path_from_id("location"), index=j))
                    trans_curves[j].keyframe_points.add(count=len(usedLocFrames[j]))
                rot_curves = []
                for j in range(4):
                    rot_curves.append(animAction.fcurves.new(poseBone.path_from_id("rotation_quaternion"), index=j))
                    rot_curves[j].keyframe_points.add(count=len(usedRotFrames[j]))
                for channelIndex in range(3):
                    totalKeyFrameInserts += len(usedLocFrames[channelIndex])
                    for j in range(len(usedLocFrames[channelIndex])):
                        frameIndex = usedLocFrames[channelIndex][j]
                        trans_curves[channelIndex].keyframe_points[j].co = (frameIndex * fpsScale, boneTrans[frameIndex][channelIndex])
                        trans_curves[channelIndex].keyframe_points[j].interpolation = 'LINEAR'
                    #end for
                #end for
                for channelIndex in range(4):
                    totalKeyFrameInserts += len(usedRotFrames[channelIndex])
                    for j in range(len(usedRotFrames[channelIndex])):
                        frameIndex = usedRotFrames[channelIndex][j]
                        rot_curves[channelIndex].keyframe_points[j].co = (frameIndex * fpsScale, boneRots[frameIndex][channelIndex])
                        rot_curves[channelIndex].keyframe_points[j].interpolation = 'LINEAR'
                    #end for
                #end for
                
                for transCurve in trans_curves:
                    transCurve.update()
                for rotCurve in rot_curves:
                    rotCurve.update()
            #end for
            frameCount = (anim["numFrames"]-1) * fpsScale
            if frameCount > maxFrames:
                maxFrames = frameCount
                
            #stash action into nla track strip
            if stashActions:
                track = animData.nla_tracks.new();
                track.name = "[Action Stash]%s" % (animAction.name)
                track.lock = True
                track.mute = True
                strip = track.strips.new(animAction.name, 0, animAction)
                strip.mute = True
            #end if
            
            progress.step()
        #end for
        
        if firstAction is not None:
            animData.action = firstAction
        
        #Exit Armature pose mode
        kex_utils.mode_set('OBJECT')
        
        if maxFrames > 1:
            bpy.context.scene.frame_end = maxFrames
        bpy.context.scene.render.fps = kex_utils.turok1FPS * fpsScale
        bpy.context.scene.render.fps_base = 1.0
        
        progress.leave_substeps("Animation Successfully Imported! (actions:%d bones:%d keyframes:%d)" % (len(anims), boneCount, totalKeyFrameInserts))
    #end with
    return True
#end func
