from numpy import dstack
import bpy
from . import fn


class AW_OT_world_space_copy(bpy.types.Operator):
    bl_idname = "pose.world_space_copy"
    bl_label = "World Copy"
    bl_description = "Copy world space transforms. Store active bone matrix"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.mode == 'POSE' and context.active_pose_bone

    def execute(self, context):
        ### only active pose bone
        # bpy.types.ViewLayer.world_space_store = context.object.matrix_world @ context.active_pose_bone.matrix
        
        if len(context.selected_pose_bones) == 1:
            bpy.types.ViewLayer.world_space_store = context.object.matrix_world @ context.active_pose_bone.matrix
            self.report({'INFO'}, 'Copied active pose bone matrix')
            return {"FINISHED"}

        pose_bones = [b for b in context.selected_pose_bones]
        armatures = set([b.id_data for b in pose_bones])
        
        pose_l = [] # pose_d = {}
        if len(armatures) == 1:
            # only relate to bone name
            for b in context.selected_pose_bones:
                # pose_d[b.name] = context.object.matrix_world @ b.matrix
                pose_l.append([b.name, b.id_data.matrix_world @ b.matrix])
            self.report({'INFO'}, f'Copied {len(pose_l)} pose bones matrices')
        
        else:
            # for obj in context.objects_in_mode_unique_data:
            for b in context.selected_pose_bones:
                # pose_d[f'{b.id_data.name}/{b.name}'] = b.id_data.matrix_world @ b.matrix
                pose_l.append([b.name, b.id_data.matrix_world @ b.matrix, b.id_data.name])
            self.report({'INFO'}, f'Copied {len(pose_l)} pose bones matrices (multi-armatures)')

        bpy.types.ViewLayer.world_space_store = pose_l

        return {"FINISHED"}

def set_matrix_and_key(pose, context=None, key=True):
    context = context or bpy.context
    ct = 0
    if isinstance(context.view_layer.world_space_store, list):    
        ## list structure : [0 bone_name, 1 matrix, 2[:armature_obj_name]]
        # Paste in parental hierarchical relation (seems right order by default but need update between bones !)
        for pb in context.selected_pose_bones:
            for blist in pose:
                
                if pb.name != blist[0]:
                    continue
                
                # skip if there is an object name and obj name
                if len(blist) > 2 and blist[2] != pb.id_data.name:
                    continue
                
                print(f'Paste "{pb.name}" position')
                pb.matrix = pb.id_data.matrix_world.inverted() @ blist[1]
                ct += 1
                
                context.evaluated_depsgraph_get() # Refresh depsgraph so children position is placed according to parent position
                break # go to next pose bone
    else:
        context.active_pose_bone.matrix = context.object.matrix_world.inverted() @ pose
        ct = 1

    if not ct:
        return

    if key:
        bpy.ops.anim.keyframe_insert(type='LocRotScale')

    return ct
            

class AW_OT_world_space_paste(bpy.types.Operator):
    bl_idname = "pose.world_space_paste"
    bl_label = "World Paste"
    bl_description = "Paste world space transforms. Apply stored matrix to active bone and key LocRotScale"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        if not hasattr(bpy.context.view_layer, 'world_space_store'):
            cls.poll_message_set("Nothing To Paste")
            return False
        return context.object and context.mode == 'POSE' and context.active_pose_bone


    def execute(self, context):
        ct = set_matrix_and_key(context.view_layer.world_space_store, context=context)
        if not ct:
            self.report({'ERROR'}, 'No pose pasted')
            return {"CANCELLED"}
        if ct != len(context.selected_pose_bones):
            self.report({'WARNING'}, f'{ct}/{len(context.selected_pose_bones)} bone position pasted')
        else:
            self.report({'INFO'}, f'{ct} bone position pasted')

        ## only if autokey is On
        # if context.scene.tool_settings.use_keyframe_insert_auto:
        #     bpy.ops.anim.keyframe_insert_menu(type='LocRotScale') # Available

        return {"FINISHED"}

class AW_OT_world_space_paste_next(bpy.types.Operator):
    bl_idname = "pose.world_space_paste_next"
    bl_label = "World Paste Jump"
    bl_description = "Paste world space transforms and keyframe available chanels\
        \nThen jump to prev/next key (on active pose bone)\
        \nKey Loc rot scale only"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        if not hasattr(bpy.context.view_layer, 'world_space_store'):
            cls.poll_message_set("Nothing To Paste")
            return False
        return context.object and context.mode == 'POSE' and context.active_pose_bone

    prev: bpy.props.BoolProperty()

    def execute(self, context):
        ct = set_matrix_and_key(context.view_layer.world_space_store, context=context)
        if not ct:
            self.report({'ERROR'}, 'No pose pasted')
            return {"CANCELLED"}
        
        # jump to next key
        act = fn.get_obj_action(context.object)
        if not act:
            self.report({'ERROR'}, 'No action on armature')
            return {'CANCELLED'}

        kx = [k.co.x for fcu in act.fcurves
        if fcu.data_path.split('"')[1] == context.active_pose_bone.bone.name
        for k in fcu.keyframe_points]

        if not kx:
            self.report({'ERROR'}, 'No keys on action available (no keyframe added)')
            return {'CANCELLED'}
        
        # for fcu in act.fcurves:
        #     if fcu.data_path.split('"')[1] == context.active_pose_bone.bone.name:

        if self.prev:
            new_frame = next((k for k in reversed(kx) if k < context.scene.frame_current), None)
        else:
            new_frame = next((k for k in kx if k > context.scene.frame_current), None)

        if not new_frame:
            self.report({'WARNING'}, 'No next frame to jump on')
            return {'FINISHED'}

        context.scene.frame_current = int(new_frame)
        return {"FINISHED"}

class AW_OT_world_space_paste_next_frame(bpy.types.Operator):
    bl_idname = "pose.world_space_paste_next_frame"
    bl_label = "World Paste Jump Frame"
    bl_description = "Paste world space transforms and keyframe available chanels\
        \nThen jump to prev/next frame\
        \nKey Loc rot scale only"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        if not hasattr(bpy.context.view_layer, 'world_space_store'):
            cls.poll_message_set("Nothing To Paste")
            return False
        return context.object and context.mode == 'POSE' and context.active_pose_bone

    prev: bpy.props.BoolProperty()

    def execute(self, context):
        ct = set_matrix_and_key(context.view_layer.world_space_store, context=context)
        if not ct:
            self.report({'ERROR'}, 'No pose pasted')
            return {"CANCELLED"}
        if ct != len(context.selected_pose_bones):
            self.report({'WARNING'}, f'{ct}/{len(context.selected_pose_bones)} bone position pasted')
        else:
            self.report({'INFO'}, f'{ct} bone position pasted')
        
        # context.object.keyframe_insert(data_path, index=-1, frame=bpy.context.scene.frame_current, group="", options={'INSERTKEY_AVAILABLE'})

        ## bpy.ops.anim.keyframe_insert(type='DEFAULT')

        ## possible type : 
        # 'Location', 'Rotation', 'Scaling', 'BUILTIN_KSI_LocRot', 'LocRotScale', 'LocRotScaleCProp',
        # 'BUILTIN_KSI_LocScale', 'BUILTIN_KSI_RotScale', 
        # 'BUILTIN_KSI_DeltaLocation', 'BUILTIN_KSI_DeltaRotation', 'BUILTIN_KSI_DeltaScale', 
        # 'BUILTIN_KSI_VisualLoc', 'BUILTIN_KSI_VisualRot', 'BUILTIN_KSI_VisualScaling', 
        # 'BUILTIN_KSI_VisualLocRot', 'BUILTIN_KSI_VisualLocRotScale', 'BUILTIN_KSI_VisualLocScale', 'BUILTIN_KSI_VisualRotScale'
        
        ## insert keyframe at value
        ## Insert Keyframes for specified Keying Set, with menu of available Keying Sets if undefined
        # bpy.ops.anim.keyframe_insert_menu(type='Available', always_prompt=False)
        
        ## Insert keyframes on the current frame for all properties in the specified Keying Set

        # jump to next frame
        offset = -1 if self.prev else 1
        new_frame = context.scene.frame_current + offset

        context.scene.frame_current = new_frame
        return {"FINISHED"}


classes=(
AW_OT_world_space_copy,
AW_OT_world_space_paste,
AW_OT_world_space_paste_next,
AW_OT_world_space_paste_next_frame,
)

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

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