import bpy
from . import fn

def get_active_nla_strip(all_nla=False):
    '''Return object active strip
    :all_nla: return first active strip found on all objects > NLA tracks
    '''

    if all_nla:
        objs = [o for o in bpy.data.objects if ob.animation_data]
    else:
        if not bpy.context.object:
            return
        objs = [bpy.context.object]
    
    for ob in objs:
        if not (anim := ob.animation_data):
            continue
        for nla in anim.nla_tracks:
            for strip in nla.strips:
                if strip.active:
                    print(f'{strip.name} on Track {nla.name}')
                    return strip

class AW_OT_nla_key_speed(bpy.types.Operator):
    bl_idname = "autowalk.nla_key_speed"
    bl_label = "NLA Key Speed"
    bl_description = "Activate animate strip time and Keyframe linear for first and last animation frame"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        nla_strip = fn.get_nla_strip(context.object) # get_active_nla_strip()
        if not nla_strip:
            self.report({'ERROR'}, 'no active NLA strip')
            return {"CANCELLED"}

        fcu = nla_strip.fcurves.find('strip_time')
        # if fcu:
        #     for k in reversed(fcu.keyframe_points):
        #         fcu.keyframe_points.remove(k)
        
        nla_strip.use_animated_time = True
        if not fcu or len(fcu.keyframe_points) == 0:
            # create if not exists
            nla_strip.strip_time = nla_strip.action_frame_start
            nla_strip.keyframe_insert('strip_time', frame=nla_strip.frame_start)

            nla_strip.strip_time = nla_strip.action_frame_end
            nla_strip.keyframe_insert('strip_time', frame=nla_strip.frame_end)

            fcu = nla_strip.fcurves.find('strip_time')
            
            # Go linear
            for k in fcu.keyframe_points:
                k.interpolation = 'LINEAR'
            
            return {"FINISHED"}

        ## if already exists : match offset (usefull when moving a strip)
        first = min([k.co.x for k in fcu.keyframe_points])
        
        offset = nla_strip.frame_start - first
        
        for k in fcu.keyframe_points:
            k.co.x += offset

        return {"FINISHED"}

class AW_OT_nla_remove_key_speed(bpy.types.Operator):
    bl_idname = "autowalk.nla_remove_key_speed"
    bl_label = "NLA Remove Key Speed"
    bl_description = "Remove strip time animation on active nla strip"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        nla_strip = fn.get_nla_strip(context.object) # get_active_nla_strip()
        if not nla_strip:
            self.report({'ERROR'}, 'no active NLA strip')
            return {"CANCELLED"}

        fcu = nla_strip.fcurves.find('strip_time')
        if fcu:
            for k in reversed(fcu.keyframe_points):
                fcu.keyframe_points.remove(k)
            nla_strip.use_animated_time = False
        else:    
            self.report({'ERROR'}, f'No strip time animation on active strip {nla_strip.name}')
            return {"CANCELLED"}

        return {"FINISHED"}

'''
class AW_OT_nla_reset_key_speed(bpy.types.Operator):
    bl_idname = "autowalk.nla_reset_key_speed"
    bl_label = "NLA Reset Key Speed"
    bl_description = "Activate animate strip time and Keyframe linear for first and last animation frame"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        nla_strip = fn.get_nla_strip(context.object) # get_active_nla_strip()
        if not nla_strip:
            self.report({'ERROR'}, 'no active NLA strip')
            return {"CANCELLED"}

        # Clear if exists (or just move first and last point ?)
        fcu = nla_strip.fcurves.find('strip_time')
        if fcu:
            for k in reversed(fcu.keyframe_points):
                fcu.keyframe_points.remove(k)
        
        nla_strip.use_animated_time = True
        nla_strip.strip_time = nla_strip.action_frame_start
        nla_strip.keyframe_insert('strip_time', frame=nla_strip.frame_start)

        nla_strip.strip_time = nla_strip.action_frame_end
        nla_strip.keyframe_insert('strip_time', frame=nla_strip.frame_end)

        fcu = nla_strip.fcurves.find('strip_time')
        if not fcu:
            return {"CANCELLED"}
        
        # Go linear
        for k in fcu.keyframe_points:
            k.interpolation = 'LINEAR'

        return {"FINISHED"}
'''

classes=(
AW_OT_nla_key_speed,
AW_OT_nla_remove_key_speed,
# AW_OT_nla_reset_key_speed,
)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)