import bpy
import mathutils


def get_armature_data(context):
   try:
       arm_object = context.active_object
       if arm_object and arm_object.type == "ARMATURE":
           return arm_object, getattr(arm_object, "pose_library", None)
       elif arm_object and arm_object.parent and arm_object.parent.type == "ARMATURE":
           return arm_object.parent, getattr(arm_object.parent, "pose_library", None)
       else:
           return None, None
   except:
       return None, None


def get_armature_action(context):
    try:
        arm_object, pose_library = get_armature_data(context)
        if getattr(arm_object.animation_data.action, "pose_markers", None):
            return getattr(arm_object.animation_data, "action", None)
    except:
        pass


def search_pose_marker(context, posename, type):
    try:
        arm_object, pose_library = get_armature_data(context)
        if type == "marker":
            return pose_library.pose_markers.get(posename, None)
        if type == "frame":
            return pose_library.pose_markers.get(posename, None).frame
        if type == "index":
            return pose_library.pose_markers.find(posename)
    except:
        pass


def find_fcurve(context, bone_name, transform, index_int):
    arm_object, pose_library = get_armature_data(context)
    pose_markers = pose_library.pose_markers
    active_frame = pose_markers.active.frame
    
    fcurve_object = pose_library.fcurves.find(
        'pose.bones["'+bone_name+'"].'+transform+'', index=index_int)
    if hasattr(fcurve_object, 'evaluate'):
        return fcurve_object.evaluate(active_frame)
    else:
        return None


def create_fcurve(context, bone_name, transform, index_int):
    arm_object, pose_library = get_armature_data(context)
    pose_markers = pose_library.pose_markers

    try:
        pose_library.fcurves.new(
            'pose.bones["'+bone_name+'"].'+transform+'', index=index_int, action_group=bone_name)
        return
    except:
        pass


def create_keyframe(context, bone_name, transform, index_int, new_marker, loc):
    arm_object, pose_library = get_armature_data(context)
    pose_markers = pose_library.pose_markers

    try:
        pose_library.fcurves.find(
            'pose.bones["'+bone_name+'"].'+transform+'', index=index_int).keyframe_points.insert(new_marker, loc)
        return
    except:
        pass


def set_keyframes_from_bones(context, arm_object, new_marker):
    none_selected = True
    for bone in arm_object.pose.bones:
        if bone.bone.select:
            none_selected = False

    for bone in arm_object.pose.bones:
        if bone.bone.select or none_selected == True:
            bone_name = bone.name

            if bone.rotation_mode == "XYZ":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "YZX":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "ZXY":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "QUATERNION":
                rot_mode = "rotation_quaternion"
            else:
                self.report({'WARNING'}, "YAPL: Unsupported bone: " + bone.name + ": " + bone.rotation_mode)
                rot_mode = None

            loc_x = bone.location[0]
            loc_y = bone.location[1]
            loc_z = bone.location[2]
            create_fcurve(context, bone_name, "location", 0)
            create_fcurve(context, bone_name, "location", 1)
            create_fcurve(context, bone_name, "location", 2)
            create_keyframe(context, bone_name, "location", 0, new_marker, loc_x)
            create_keyframe(context, bone_name, "location", 1, new_marker, loc_y)
            create_keyframe(context, bone_name, "location", 2, new_marker, loc_z)
            if rot_mode == "rotation_quaternion":
                rot_w = bone.rotation_quaternion[0]
                rot_x = bone.rotation_quaternion[1]
                rot_y = bone.rotation_quaternion[2]
                rot_z = bone.rotation_quaternion[3]
                create_fcurve(context, bone_name, rot_mode, 0)
                create_fcurve(context, bone_name, rot_mode, 1)
                create_fcurve(context, bone_name, rot_mode, 2)
                create_fcurve(context, bone_name, rot_mode, 3)
                create_keyframe(context, bone_name, rot_mode, 0, new_marker, rot_w)
                create_keyframe(context, bone_name, rot_mode, 1, new_marker, rot_x)
                create_keyframe(context, bone_name, rot_mode, 2, new_marker, rot_y)
                create_keyframe(context, bone_name, rot_mode, 3, new_marker, rot_z)
            elif rot_mode == "rotation_euler":
                rot_x = bone.rotation_euler[0]
                rot_y = bone.rotation_euler[1]
                rot_z = bone.rotation_euler[2]
                create_fcurve(context, bone_name, rot_mode, 0)
                create_fcurve(context, bone_name, rot_mode, 1)
                create_fcurve(context, bone_name, rot_mode, 2)
                create_keyframe(context, bone_name, rot_mode, 0, new_marker, rot_x)
                create_keyframe(context, bone_name, rot_mode, 1, new_marker, rot_y)
                create_keyframe(context, bone_name, rot_mode, 2, new_marker, rot_z)
            scl_x = bone.scale[0]
            scl_y = bone.scale[1]
            scl_z = bone.scale[2]
            create_fcurve(context, bone_name, "scale", 0)
            create_fcurve(context, bone_name, "scale", 1)
            create_fcurve(context, bone_name, "scale", 2)
            create_keyframe(context, bone_name, "scale", 0, new_marker, scl_x)
            create_keyframe(context, bone_name, "scale", 1, new_marker, scl_y)
            create_keyframe(context, bone_name, "scale", 2, new_marker, scl_z)


def set_bones_from_keyframes(context, arm_object, active_marker):
    none_selected = True
    for bone in arm_object.pose.bones:
        if bone.bone.select:
            none_selected = False

    for bone in arm_object.pose.bones:
        if bone.bone.select or none_selected == True:
            bone_name = bone.name

            if bone.rotation_mode == "XYZ":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "YZX":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "ZXY":
                rot_mode = "rotation_euler"
            elif bone.rotation_mode == "QUATERNION":
                rot_mode = "rotation_quaternion"
            else:
                self.report({'WARNING'}, "YAPL: Unsupported bone: " + bone.name + ": " + bone.rotation_mode)
                rot_mode = None

            loc_x = find_fcurve(context, bone_name, "location", 0) or 0.0
            loc_y = find_fcurve(context, bone_name, "location", 1) or 0.0
            loc_z = find_fcurve(context, bone_name, "location", 2) or 0.0
            if rot_mode == "rotation_quaternion":
                rot_w = find_fcurve(context, bone_name, rot_mode, 0) or 1.0
                rot_x = find_fcurve(context, bone_name, rot_mode, 1) or 0.0
                rot_y = find_fcurve(context, bone_name, rot_mode, 2) or 0.0
                rot_z = find_fcurve(context, bone_name, rot_mode, 3) or 0.0
            elif rot_mode == "rotation_euler":
                rot_x = find_fcurve(context, bone_name, rot_mode, 0) or 0.0
                rot_y = find_fcurve(context, bone_name, rot_mode, 1) or 0.0
                rot_z = find_fcurve(context, bone_name, rot_mode, 2) or 0.0
            scl_x = find_fcurve(context, bone_name, "scale", 0) or 1.0
            scl_y = find_fcurve(context, bone_name, "scale", 1) or 1.0
            scl_z = find_fcurve(context, bone_name, "scale", 2) or 1.0

            bone.location = mathutils.Vector((loc_x, loc_y, loc_z))
            if bone.rotation_mode == "XYZ":
                bone.rotation_euler = mathutils.Euler(
                    (rot_x, rot_y, rot_z))
            elif bone.rotation_mode == "YZX":
                bone.rotation_euler = mathutils.Euler(
                    (rot_x, rot_y, rot_z))
            elif bone.rotation_mode == "ZXY":
                bone.rotation_euler = mathutils.Euler(
                    (rot_z, rot_x, rot_y))
            elif bone.rotation_mode == "YXZ":
                bone.rotation_euler = mathutils.Euler(
                    (rot_y, rot_x, rot_z))
            elif bone.rotation_mode == "XZY":
                bone.rotation_euler = mathutils.Euler(
                    (rot_x, rot_z, rot_y))
            elif rot_mode == "rotation_quaternion":
                bone.rotation_quaternion = mathutils.Quaternion(
                    (rot_w, rot_x, rot_y, rot_z))
            bone.scale = mathutils.Vector((scl_x, scl_y, scl_z))