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 = 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,
)

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

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