#! /bin/bash
# vim: shiftwidth=4 tabstop=4 expandtab

# Default configuration values
[[ "${SCHEDULED_SNAPSHOTS_NAME_FORMAT:-null}" == "null" ]] && \
    SCHEDULED_SNAPSHOTS_NAME_FORMAT="%Y%m%d-%H%M%S-DISASTER-RECOVERY"

[[ "${SCHEDULED_SNAPSHOTS_RETENTION:-null}" == "null" ]] && SCHEDULED_SNAPSHOTS_RETENTION=1
[[ "${SCHEDULED_SNAPSHOTS_VMS_RETENTION[*]:-null}" == "null" ]] && \
    declare -A SCHEDULED_SNAPSHOTS_VMS_RETENTION=()
[[ "${SCHEDULED_SNAPSHOTS_SHARED_VOLUMES_RETENTION[*]:-null}" == "null" ]] && \
    declare -A SCHEDULED_SNAPSHOTS_SHARED_VOLUMES_RETENTION=()

[[ "${SCHEDULED_SNAPSHOTS_INCLUDED_VMS[*]:-null}" == "null" ]] && \
    SCHEDULED_SNAPSHOTS_INCLUDED_VMS=()

[[ "${SCHEDULED_SNAPSHOTS_DAYS[*]:-null}" == "null" ]] && SCHEDULED_SNAPSHOTS_DAYS=(1 2 3 4 5)
[[ "${SCHEDULED_SNAPSHOTS_VMS_DAY[*]:-null}" == "null" ]] && declare -A SCHEDULED_SNAPSHOTS_VMS_DAY
[[ "${SCHEDULED_SNAPSHOTS_VMS_VOLUMES_PATTERN[*]:-null}" == "null" ]] && \
    declare -A SCHEDULED_SNAPSHOTS_VMS_VOLUMES_PATTERN
[[ "${SCHEDULED_SNAPSHOTS_WAIT_DELAY:-null}" == "null" ]] && SCHEDULED_SNAPSHOTS_WAIT_DELAY=5

function check_weekday() {
    in_array "$1" $( seq 1 7 ) && return 0
    return 1
}

function format_weekday() {
    local day_number=$1 current_weekday day_name offset
    current_weekday=$(date +%u)
    offset=$((day_number - current_weekday))
    day_name=$( LC_ALL=C date -d "$offset days" +%A )
    echo "${day_name,,}"
}

function check_scheduled_snapshots_name_pattern() {
    # If no pattern is provided or configured, generate it from SCHEDULED_SNAPSHOTS_NAME_FORMAT
    [[ "${SCHEDULED_SNAPSHOTS_NAME_PATTERN:-null}" == "null" ]] && \
        generate_scheduled_snapshots_name_pattern

    local result
    result=$(
        jo name="$( date "+$SCHEDULED_SNAPSHOTS_NAME_FORMAT" )" |
        jq -r --arg pattern "$SCHEDULED_SNAPSHOTS_NAME_PATTERN" '.name | test($pattern)'
    )
    [[ "$result" == "true" ]] && return 0
    return 1
}

function get_scheduled_snapshots_planning() {
    local output_var=$1
    declare -Ag "$output_var=()"
    local vm day

    log INFO "List VMs to handle..."
    get_scheduled_snapshots_vms_list
    log INFO "${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]} VMs to handle"

    log INFO "Explode VMs by week days..."
    # Explode included VMs by day:
    # - exclude inactive VMs
    # - taking into account VMs explicitly configured to be managed on Friday
    # - distributed the other VMs over the other days

    for day in "${SCHEDULED_SNAPSHOTS_DAYS[@]}"; do
        declare -gA "${output_var}[$day]="
    done

    # Handle VMs with custom scheduled snapshot day
    if [[ ${#SCHEDULED_SNAPSHOTS_VMS_DAY[@]} -gt 0 ]]; then
        for vm in "${!SCHEDULED_SNAPSHOTS_VMS_DAY[@]}"; do
            day=${SCHEDULED_SNAPSHOTS_VMS_DAY[$vm]}
            if ! in_array "$vm" "${SCHEDULED_SNAPSHOTS_VMS_LIST[@]}"; then
                log ERROR "VM $vm with scheduled snapshots on $(format_weekday "$day") is not" \
                    "configured to be handled, skipping it."
                continue
            fi
            log DEBUG "Scheduled snapshots of VM $vm snapshots explicitly configured on" \
                "$(format_weekday "$day")"
            declare -gA "${output_var}[$day]+=$vm "
        done
    fi

    # Distribute other VMs over the other days

    # Compute how many VMs we have to handle by day
    (( VMS_TO_HANDLE_BY_DAY=${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]}/${#SCHEDULED_SNAPSHOTS_DAYS[@]} ))
    ((
        EXTRA_VMS_TO_HANDLE_ON_FIRST_WEEKDAY=${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]} -
            (VMS_TO_HANDLE_BY_DAY*${#SCHEDULED_SNAPSHOTS_DAYS[@]})
    ))
        log DEBUG "Need to schedule snapshot of $VMS_TO_HANDLE_BY_DAY by day on" \
        "${#SCHEDULED_SNAPSHOTS_DAYS[@]} days with one extra VM on first" \
        "$EXTRA_VMS_TO_HANDLE_ON_FIRST_WEEKDAY weekdays to handle the" \
        "${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]} VMs."

    local idx=0 day_vms
    for day in "${SCHEDULED_SNAPSHOTS_DAYS[@]}"; do
        while true; do
            day_vms=$( wc -w <<< "$(eval echo "\${${output_var}[$day]}")" )
            [[ $day_vms -gt $VMS_TO_HANDLE_BY_DAY ]] && break
            if [[ $day_vms -eq $VMS_TO_HANDLE_BY_DAY ]]; then
                if [[ $EXTRA_VMS_TO_HANDLE_ON_FIRST_WEEKDAY -gt 0 ]]; then
                    log DEBUG "Schedule one extra VM snapshot on $(format_weekday "$day")"
                    (( EXTRA_VMS_TO_HANDLE_ON_FIRST_WEEKDAY-=1 ))
                else
                    break
                fi
            fi

            vm=${SCHEDULED_SNAPSHOTS_VMS_LIST[$idx]}
            # Ignore VMs with custom scheduled snapshot day
            while [[ "${SCHEDULED_SNAPSHOTS_VMS_DAY[$vm]:-null}" != "null" ]]; do
                (( idx+=1 ))
                vm=${SCHEDULED_SNAPSHOTS_VMS_LIST[$idx]}
            done
            log DEBUG "Scheduled snapshots of VM $vm snapshots on $(format_weekday "$day")"
            declare -gA "${output_var}[$day]+=$vm "
            (( idx+=1 ))
            [[ $idx -ge ${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]} ]] && break
        done
        [[ $idx -ge ${#SCHEDULED_SNAPSHOTS_VMS_LIST[@]} ]] && break
    done
}

# Get list of VMs with scheduled snapshots: sort and deduplicate VM list in
# SCHEDULED_SNAPSHOTS_INCLUDED_VMS and store result in global array SCHEDULED_SNAPSHOTS_VMS_LIST
function get_scheduled_snapshots_vms_list() {
    [[ "${SCHEDULED_SNAPSHOTS_VMS_LIST[*]:-null}" != "null" ]] && return
    declare -ag SCHEDULED_SNAPSHOTS_VMS_LIST=()
    mapfile -t SCHEDULED_SNAPSHOTS_VMS_LIST < <(
        implode $'\n' "${SCHEDULED_SNAPSHOTS_INCLUDED_VMS[@]}" | sort -u
    )
}

# Get the snapshot retention for a given RBD volume
function get_scheduled_snapshots_retention(){
    local volume=$1
    local output_var=$2
    local verbose=$3
    if is_shared_volume "$volume" "$verbose"; then
        log DEBUG "Volume $volume is shared by multiple VMs: ${SHARED_VOLUMES[$volume]}"
        if [[ "${SCHEDULED_SNAPSHOTS_SHARED_VOLUMES_RETENTION[$volume]:-null}" != "null" ]]; then
            log DEBUG "Shared volume $volume have a custom retention policy of its scheduled" \
                "snapshots: keep ${SCHEDULED_SNAPSHOTS_SHARED_VOLUMES_RETENTION[$volume]}" \
                "snapshot(s)."
            declare -g "$output_var=${SCHEDULED_SNAPSHOTS_SHARED_VOLUMES_RETENTION[$volume]}"
        else
            log DEBUG "Shared volume $volume does not have custom retention policy, use default:"\
                "keep $SCHEDULED_SNAPSHOTS_RETENTION snapshot(s)."
            declare -g "$output_var=$SCHEDULED_SNAPSHOTS_RETENTION"
        fi
        return
    fi

    get_scheduled_snapshots_vms_list
    get_volume_vm "$volume" VOLUME_VM "$verbose"
    if [[ -z "$VOLUME_VM" ]]; then
        log WARNING "The volume '$volume' is not attached to any VM: drop all of its existing" \
            "scheduled snapshots."
        declare -g "$output_var=0"
    elif ! in_array "$VOLUME_VM" "${SCHEDULED_SNAPSHOTS_VMS_LIST[@]}"; then
        log WARNING "The volume '$volume' is attached to VM $VOLUME_VM not configured to have" \
            "scheduled snapshots: drop all of its existing scheduled snapshots."
        declare -g "$output_var=0"
    elif [[ "${SCHEDULED_SNAPSHOTS_VMS_RETENTION[$VOLUME_VM]:-null}" != "null" ]]; then
        log DEBUG "VM $VOLUME_VM have a custom retention policy of its scheduled snapshots:" \
            "keep ${SCHEDULED_SNAPSHOTS_VMS_RETENTION[$VOLUME_VM]} snapshot(s)."
        declare -g "$output_var=${SCHEDULED_SNAPSHOTS_VMS_RETENTION[$VOLUME_VM]}"
    else
        log DEBUG "VM $VOLUME_VM does not have custom retention policy, use default:"\
            "keep $SCHEDULED_SNAPSHOTS_RETENTION snapshot(s)."
        declare -g "$output_var=$SCHEDULED_SNAPSHOTS_RETENTION"
    fi
}

function generate_scheduled_snapshots_name_pattern {
    declare -g SCHEDULED_SNAPSHOTS_NAME_PATTERN
    SCHEDULED_SNAPSHOTS_NAME_PATTERN=$(
        sed -e 's/%[dHImMSy]/[0-9]{2}/g' \
            -e 's/%Y/[0-9]{4}/g' \
            -e 's/%j/[0-9]{3}/g' \
            -e 's/%q/[0-9]/g' \
            -e 's/%u/[1-7]/g' \
            -e 's/%[UW]/[0-5][0-9]/g' \
            -e 's/%w/[0-6]/g' \
            -e 's/%F/[0-9]{4}-[0-9]{2}-[0-9]{2}/g' \
            -e 's/%D/[0-9]{2}-[0-9]{2}-[0-9]{2}/g' \
            -e 's/%R/[0-9]{2}:[0-9]{2}/g' \
            -e 's/%T/[0-9]{2}:[0-9]{2}:[0-9]{2}/g' \
            -e 's/%s/[0-9]{10}/g' \
            -e 's/%z/[+-][0-1][0-9][0-6][0-9]/g' \
            -e 's/%:z/[+-][0-1][0-9]:[0-6][0-9]/g' \
            -e 's/%::z/[+-][0-1][0-9]:[0-6][0-9]:[0-6][0-9]/g' \
            -e 's/^/^/' \
            -e 's/$/$/' <<< "$SCHEDULED_SNAPSHOTS_NAME_FORMAT"
    )
    [[ $( grep -cEi '%[a-z]' <<< "$SCHEDULED_SNAPSHOTS_NAME_PATTERN" ) -ne 0 ]] && \
        log FATAL "Invalid snapshot name pattern computed from name format" \
            "($SCHEDULED_SNAPSHOTS_NAME_PATTERN). May be it contain unsupported date" \
            "replacement pattern."
}

function list_volumes_with_scheduled_snapshots() {
    local pool=$1
    local output_var=$2

    # If no pattern is provided or configured, generate it from SCHEDULED_SNAPSHOTS_NAME_FORMAT
    [[ "${SCHEDULED_SNAPSHOTS_NAME_PATTERN:-null}" == "null" ]] && \
        generate_scheduled_snapshots_name_pattern

    declare -ga "$output_var=()"
    local volume
    for volume in $(
        rbd ls -l "$pool" --format=json | \
            jq -r --arg pattern "$SCHEDULED_SNAPSHOTS_NAME_PATTERN" \
            '[ .[] | select(.snapshot? and (.snapshot | test($pattern))) | .image ]
            | unique | sort[]'
    ); do
        declare -ga "$output_var+=( '$volume' )"
    done
}

# Get scheduled snapshots info and sort them by timestamp
function get_volume_scheduled_snapshots() {
    # If no pattern is provided or configured, generate it from SCHEDULED_SNAPSHOTS_NAME_FORMAT
    [[ "${SCHEDULED_SNAPSHOTS_NAME_PATTERN:-null}" == "null" ]] && \
        generate_scheduled_snapshots_name_pattern

    local output_var=$2 snaps
    get_snapshots_info "$@" 2> >(grep -v "fast-diff map is not enabled")

    snaps=$(
        jq \
            --arg pattern "$SCHEDULED_SNAPSHOTS_NAME_PATTERN" \
            --raw-output \
            '
                [
                    .[] | select( .name | test($pattern) )
                ]
                | sort_by(.timestamp)
            ' <<< "${!output_var}"
    )
    declare -g "$output_var=$snaps"
}

# check and log ceph cluster status
# - report snaptrims status and non HEALTH_OK status
# - return 1 if currently snaptrims and/or non HEALTH_OK status, 0 otherwise
function check_ceph_snaptrims_and_health_status() {
    local ceph_status snaptrims health_status
    ceph_status=$(ceph status --format json)
    snaptrims=$(
        jq -r '
            .pgmap.pgs_by_state[] | select(.state_name | match("snaptrim")) |
            .state_name + ": " + (.count|tostring)' <<< "$ceph_status" |
        sed -z 's/\n/, /g;s/, $/\n/'
    )
    if [[ -n "$snaptrims" ]]; then
        display "snaptrim ($snaptrims)..."
        return 1
    fi

    health_status=$( jq -r '.health.status' <<< "$ceph_status" )
    if [[ "$health_status" == "HEALTH_WARN" ]]; then
        local -a checks check errors
        mapfile -t checks < <( jq -r '.health.checks | keys[]' <<< "$ceph_status" )
        for check in "${checks[@]}"; do
            case "$check" in
                OSD_NEARFULL|POOL_NEARFULL|BLUEFS_SPILLOVER)
                    display "$(
                        printf "Ignore $check warning (%s)" \
                            "$(jq -r ".health.checks.$check.summary.message" <<< "$ceph_status")"
                    )"
                    ;;
                *)
                    errors+=( "$check" )
                    ;;
            esac
        done
        if [[ "${#errors[@]}" -gt 0 ]]; then
            log WARNING "HEALTH NON OK (${#errors[@]} problem(s))"
            ceph health detail 2>&1
            return 1
        fi
    fi
    return 0
}

# Wait until the ceph cluster is not snaptrims or in non HEALTH_OK status
function wait_until_ceph_cluster_not_snaptrims_and_health_ok(){
    if [[ $1 -eq 1 ]]; then
        display "Waiting 5 seconds to ensure previous deletion start to impact cluster..."
        sleep 5
    fi
    while ! check_ceph_snaptrims_and_health_status ; do
        sleep 10
    done
}

# Cleanup one volume's scheduled snapshots
function cleanup_volume_scheduled_snapshots(){
    local volume=$1 snaps_count snaps_to_del
    get_scheduled_snapshots_retention "$volume" RETENTION

    log INFO "Get snapshots of volume $volume..."
    get_volume_scheduled_snapshots "$volume" SNAPS with-size

    if [[ -z "$SNAPS" ]]; then
        log ERROR "Failed to retrieve snapshots of volume $volume"
        return 1
    fi

    snaps_count=$( jq length <<< "$SNAPS" )
    log INFO "Found ${snaps_count} snapshots for ${volume} (configured retention: $RETENTION)"

    [[ $snaps_count -le $RETENTION ]] && return

    snaps_to_del=$(
        jq \
            --argjson retention "$RETENTION" \
            --raw-output \
            --compact-output \
            '
                [ .[] ]
                [:-$retention]
                | .[]
            ' <<< "$SNAPS"
    )

    set_newline_ifs
    local snap_data snap size start_time duration
    for snap_data in $snaps_to_del; do
        snap=$( jq -r .name <<< "$snap_data" )
        size=$( jq -r .used_size <<< "$snap_data" )

        log INFO "Removing $volume@$snap (size: $( format_size -b -z "$size" ))"

        # Unprotect the snapshot
        wait_until_ceph_cluster_not_snaptrims_and_health_ok
        run_external rbd snap unprotect "$volume@$snap" 2> >(
            grep -v 'snap is already unprotected' >&2
        )

        # Remove the snapshot
        wait_until_ceph_cluster_not_snaptrims_and_health_ok
        start_time=$(date +%s)
        run_external rbd snap rm "$volume@$snap" 2> >( grep -v 'Removing snap:.*done.' >&2 )

        # Log duration
        wait_until_ceph_cluster_not_snaptrims_and_health_ok 1
        (( duration=$(date +%s)-start_time ))
        log INFO "Snapshot $volume@$snap (size: $( format_size -b -z "$size" ))" \
            "removed in $(format_duration $duration)"
    done
    restore_ifs
}
