import bpy


def getArmatureObject(context):
    obj = context.active_object
    if obj:
        if obj.type == "ARMATURE":
            return obj
        elif obj.parent is not None and obj.parent.type == "ARMATURE":
            return obj.parent
        else:
            return None


def getLegacyPoseLibrary(context):
    arm = getArmatureObject(context)
    return getattr(arm, "pose_library", None)


def getArmatureAction(context):
    arm = getArmatureObject(context)
    return getattr(arm.animation_data, "action", None)


def getDsplAction(context):
    arm = getArmatureObject(context)
    return getattr(arm.dspl, "pose_library", None)


def getPoseLib(context):
    arm = getArmatureObject(context)
    pose_library_dspl = getDsplAction(context)
    pose_library_action = getArmatureAction(context)
    pose_library_legacy = getLegacyPoseLibrary(context)

    if pose_library_dspl:
        return pose_library_dspl
    elif pose_library_action and not pose_library_dspl:
        return pose_library_action
    elif pose_library_legacy and not pose_library_dspl:
        return pose_library_legacy
    else:
        return None


def searchPoseMarker(context, posename, type):
    try:
        pose_library = getPoseLib(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 findFcurve(context, bone_name, transform, index_int):
    arm_object = getArmatureObject(context)
    action_object = getPoseLib(context)
    pose_markers = action_object.pose_markers
    active_frame = pose_markers.active.frame
    
    fcurve_object = action_object.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 createFcurve(context, bone_name, transform, index_int):
    arm_object = getArmatureObject(context)
    action_object = getPoseLib(context)
    pose_markers = action_object.pose_markers

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


def createKeyframe(context, bone_name, transform, index_int, new_marker, loc):
    arm_object = getArmatureObject(context)
    action_object = getPoseLib(context)
    pose_markers = action_object.pose_markers

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


def setKeyframesFromBones(context, new_marker):
    arm_object = getArmatureObject(context)
    for bone in arm_object.pose.bones:
        if bone.bone.select or arm_object.dsplvars.only_selected == False:
            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'}, "DSPL: 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]
            createFcurve(context, bone_name, "location", 0)
            createFcurve(context, bone_name, "location", 1)
            createFcurve(context, bone_name, "location", 2)
            createKeyframe(context, bone_name, "location", 0, new_marker, loc_x)
            createKeyframe(context, bone_name, "location", 1, new_marker, loc_y)
            createKeyframe(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]
                createFcurve(context, bone_name, rot_mode, 0)
                createFcurve(context, bone_name, rot_mode, 1)
                createFcurve(context, bone_name, rot_mode, 2)
                createFcurve(context, bone_name, rot_mode, 3)
                createKeyframe(context, bone_name, rot_mode, 0, new_marker, rot_w)
                createKeyframe(context, bone_name, rot_mode, 1, new_marker, rot_x)
                createKeyframe(context, bone_name, rot_mode, 2, new_marker, rot_y)
                createKeyframe(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]
                createFcurve(context, bone_name, rot_mode, 0)
                createFcurve(context, bone_name, rot_mode, 1)
                createFcurve(context, bone_name, rot_mode, 2)
                createKeyframe(context, bone_name, rot_mode, 0, new_marker, rot_x)
                createKeyframe(context, bone_name, rot_mode, 1, new_marker, rot_y)
                createKeyframe(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]
            createFcurve(context, bone_name, "scale", 0)
            createFcurve(context, bone_name, "scale", 1)
            createFcurve(context, bone_name, "scale", 2)
            createKeyframe(context, bone_name, "scale", 0, new_marker, scl_x)
            createKeyframe(context, bone_name, "scale", 1, new_marker, scl_y)
            createKeyframe(context, bone_name, "scale", 2, new_marker, scl_z)