#!/bin/bash

CONFIG="/etc/squashfu.conf"
source "$CONFIG"

# Informational output w/ happy colors
debug () {
    if [[ "$DEBUG" == "true" ]]; then
        printf '\033[1;33mDEBUG ::\033[1;m %s\n' "$*"
    fi
}

info () {
    printf '\033[1;34m::\033[1;m %s\n' "$*"
}

die () {
    printf '\033[1;31mFATAL ::\033[1;m %s\n' "$*" >&2
    exit 1
}

create_new_squash () {
# Args: number of bins to be squashed (as determined by check_for_resquash), -1 on initial creation
# Returns: 0 on success, non-zero on failure

    # If making first seed, create it empty and return
    if [[ $1 -eq -1 ]]; then
        mksquashfs "$UNION_MOUNT" "$SEED" -b 65536
        return $?
    fi

    # Determine oldest $1 bins and mount them with the current squash
    local old_bins=($(sort -n -r -t: -k2 "$BINVENTORY" | tail -$1 | cut -d: -f1))

    debug "old_bins declared as: ${old_bins[@]}"

    mount_union_with_bins ${old_bins[@]}

    # Create new squash with temp name
    mksquashfs "$UNION_MOUNT" "$SEED.replace" -b 65536

    # If the squash wasn't made correctly, we don't want to continue
    if [[ $? -ne 0 ]]; then
        return 1
    fi

    unmount_all

    # Replace old squash
    mv "${SEED}.replace" "$SEED"

    # Delete old bins, and remove entry from binventory
    for bin in ${old_bins[@]}; do
        rm -rf "${BINS_DIR}/$bin"
        sed -i "/^$bin:/d" "$BINVENTORY"
    done

    # Clean up $binventory
    sweep_bins
}

create_new_incremental () {
# Args: none
# Returns: 0 on success, non-zero on error

    # Make a new bin for this incremenetal
    get_next_available_bin
    create_new_bin $?

    # Determine the mount order via binventory
    local bin_order=($(sort -n -r -t: -k2 "$BINVENTORY" | cut -d: -f1))

    mount_squash
    mount_union_with_bins ${bin_order[@]}

    # Die with error on mount, else start rsync
    if [[ $? -ne 0 ]]; then
        return 1;
    fi

    # Includes should be pulled in directly from config and not need doctoring
    #INCLUDES=($(sed -n '/^<<INCLUDES$/,/^INCLUDES$/p' $CONFIG | grep -vE "^<*INCLUDES$"))
    #EXCLUDES=($(sed -n '/^<<EXCLUDES$/,/^EXCLUDES$/p' $CONFIG | \
    #            grep -vE "^<*EXCLUDES$" | \
    #            sed -n 's/\(.*\)/--exclude \1/p'))
    EXCLUDES=$(for excl in $EXCLUDES; do echo --exclude $excl; done)
    exit 1

    debug "rsync ${RSYNC_OPTS[@]} ${INCLUDES[@]} ${EXCLUDES[@]} "$UNION_MOUNT""
    /usr/bin/rsync ${RSYNC_OPTS[@]} ${INCLUDES[@]} ${EXCLUDES[@]} "$UNION_MOUNT"

    return $?
}

create_new_bin () {
# Arguments: 1, the number of the bin to create
# Returns: 0 on success, non-zero on error

    debug "Asked to create new bin: $1"
    # Create new directory, fail if it exists (something's wrong)
    mkdir "${BINS_DIR}/$1"
    if [[ $? -ne 0 ]]; then
        return $?
    fi

    # Update binventory with new bin name and timestamp
    debug "Updating $BINVENTORY, adding bin $1"
    echo "${1}:$(date +%s)" >> "$BINVENTORY"

    # If write to bin list fails, remove diretory and exit
    if [[ $? -ne 0 ]]; then
        rmdir "${BINS_DIR}/${1}"
        die "Error writing to '$BINVENTORY'"
    fi

    return
}

# Mounting functions
mount_squash () {
# Arguments: none
# Returns: return code of mount command
    debug "Mounting Squash"
    mount -o loop,ro "$SEED" "$SQUASH_MOUNT"
    return $?
}

mount_union_with_bins () {
# Arguments: numbers of bins to be mounted (variable number)
# Returns: 0 on successful mount, non-zero on failure
    debug "Requested to mount bins: $*"

    # Mount first as rw, shift, and mount the rest ro
    branches="br=${BINS_DIR}/$1=rw:"; shift
    if [[ -n $1 ]]; then
        for bin in $*; do
            branches="${branches}${BINS_DIR}/$bin=ro:"
        done
    fi
    branches="${branches}${SQUASH_MOUNT}=ro"

    debug "mount -t aufs none "$UNION_MOUNT" -o udba=reval,$branches"
    mount -t aufs none "$UNION_MOUNT" -o udba=reval,$branches

    return $?
}

# Unmounting functions
unmount_union () {
# Args: none
# Returns: return code from umount
    debug "Unmounting union"
    while [[ $(mountpoint "$UNION_MOUNT" | grep "is a mount") ]]; do
        umount "$UNION_MOUNT" 2>/dev/null
        sleep 1
    done
    return $?
}

unmount_squash () {
# Args: none
# Returns: return code from umount
    debug "Unmounting squash"
    while [[ $(mountpoint "$SQUASH_MOUNT" | grep "is a mount") ]]; do
        umount "$SQUASH_MOUNT" 2>/dev/null
        sleep 1
    done
    return $?
}

unmount_all () {
# Args: none
# Returns: none

    if [[ $UID -ne 0 ]]; then
        die "Must be root to unmount."
    fi

    # Union MUST be unmounted first
    unmount_union
    unmount_squash
}

check_for_resquash () {
# Args: none
# Returns: number of bins needing to be merged
    local number_of_bins=$(grep -vE "^[ ]*$" "$BINVENTORY" | wc -l)
    debug "Found $number_of_bins bins"

    if [[ $number_of_bins -gt $MAX_BINS ]]; then
        return $[ $number_of_bins - $MIN_BINS ]
    else
        return 0
    fi
}

get_next_available_bin () {
# Arguments: none
# Returns: Numeric value of the next unused bin
    next_bin=$[ $(cut -d: -f1 "$BINVENTORY" | sort -n | tail -1) + 1 ]
    debug "Next available bin = $next_bin"
    return $next_bin
}

sweep_bins () {
# Arguments: none
# Returns: none
    debug "Entering sweep_bins"
    count=1

    # Make sure bins are numbered in order, clean up if not. In other words,
    # if we have 10 bins, make sure they're ordered 1 through 10.
    ls "${BINS_DIR}" | while read bin; do
    #for bin in "${BINS_DIR}/*"; do
        if [[ ! -d "${BINS_DIR}/$count" ]]; then
            high_bin=$(ls "${BINS_DIR}" | sort -n | tail -1)
            mv "${BINS_DIR}/$high_bin" "${BINS_DIR}/$count"
            sed -i "/^$high_bin:/s/^$high_bin:/$count:/" "$BINVENTORY"
        fi
        count=$[ $count + 1 ]
    done
    debug "Leaving sweep_bins"
    ls "$BINS_DIR"

}

action_backup () {
# Args: options array squashfu was invoked with, shifted 1
# Returns: none

    if [[ $UID -ne 0 ]]; then
        die "Must be root to perform a backup"
    fi

    # Does the binventory exist? If not, prompt to make sure this is an initialization
    # FIRST_RUN=0
    if [[ ! -f "$BINVENTORY" || ! -f "$SEED" ]]; then
        read -p "Looks like this is your first time running SquashFu. Is this correct? (y/n) " ans
        while [[ true ]]; do
            case $ans in
                [yY]) break ;;
                [nN]) die "Your bin inventory and/or seed seem to be missing. Please fix this before continuing." ;;
                *) ;;
            esac
        done

        # If we got here, the user answered yes, so initialize a new structure
        mkdir -p "$UNION_MOUNT"
        mkdir -p "$SQUASH_MOUNT"
        mkdir -p "${BINS_DIR}"
        touch "$BINVENTORY"
        create_new_squash -1
        FIRST_RUN=1
    fi

    # Cleanup mounts, in case user was doing a rollback and forgot to unmount (or error on last run)
    unmount_all

    create_new_incremental

    check_for_resquash
    if [[ val=$? -gt 0 ]]; then
        create_new_squash $val
    elif [[ $FIRST_RUN -eq 1 ]]; then
        create_new_squash 1
    fi

    # TODO: Report if requested

    unmount_all
}

action_rollback () {
# Args: number of backups to roll back
# Returns: none

    if [[ $UID -ne 0 ]]; then
        die "Must be root to perform a rollback"
    fi

    # Validate input with test cases
    if [[ -z $1 ]]; then
        die "The rollback action requires 1 additional argument."
    fi

    if [[ $1 -le 0 ]]; then
        die "Please provide a positive number of backups to roll back"
    fi

    # Form a chronologically ordered list of bins, assuming the user didn't give bogus input
    local bin_list=($(grep -vE "^[ ]*$" "$BINVENTORY" | sort -t: -r -n -k2 | cut -d: -f1))

    if [[ $1 -gt ${#bin_list[@]} ]]; then
        die "Cannot rollback more than ${#bin_list[@]} backups"
    fi

    local num_to_mount=$[ ${#bin_list[@]} - $1 ]

    mountpoint "$UNION_MOUNT" || mountpoint "$SQUASH_MOUNT" && unmount_all

    mount_squash

    mount_union_with_bins ${bin_list[@]:(-$num_to_mount)}

    info "Union is now mounted at '$UNION_MOUNT'"
}

action_report () {
    info "SquashFu Usage Report"
    echo
    # Enumerate bins, sort date order, print human readable create date and size
    OLDIFS=$IFS;IFS='$:'
    printf "%10s\t%25s\t%7s\n" "Bin" "Date Created" "Size"
    grep -vE "^[\t ]*$" "$BINVENTORY" | sort -r -k2 -n | while read bin stamp; do
        printf "%10d\t%25s\t%7s\n" $bin \
               "$(date --rfc-3339=seconds --date="1970-01-01 $stamp sec GMT")" \
               "$(du -sh ${BINS_DIR}/$bin 2>/dev/null | awk '{print $1}')"
    done
    IFS=$OLDIFS
    printf "%10s\t%25s\t%7s\n" "" "Incremental Total" "$(du -sh "$BINS_DIR" 2>/dev/null | awk '{print $1}')"

    # Print totals (not efficient -- reruns du on things we already ran it on)
    printf "\n%10s\t%25s\t%7s\n" "" "$(basename $SEED)" "$(du -h "$SEED" 2>/dev/null | awk '{print $1}')"
    printf "\n%10s\t%25s\t%7s\n" "" "Grand Total" \
           "$(du -csh "$BINS_DIR" "$SEED" 2>/dev/null | grep -E "total$" | awk '{print $1}')"
}

usage () {
    info "SquashFu: Super Awesome Backup Express (Professional Edition)"
    cat <<HELP

USAGE
    squashfu <operation>

OPERATIONS
  -B
    Runs a regular backup, using the config file at /etc/squashfu, unless
    otherwise specified with the -c option.

  -Q
    Displays the size of the seed, the incrementals, and the actual backup. If
    you provide no additional options, a basic report will be given. Specifying
    "full" will give more detail about individual bins.

  -R <number of bins>
    Rollback specified number of backups and mount union for browsing. The rolled
    back data will be mounted at $UNION_MOUNT.

  -U
    Unmount squash and union. Although SquashFu will always check and unmount as 
    necessary before an operation, this is provided as a safeguard.

HELP
    exit 0
}

case $1 in
    "-B") action_backup ;;
    "-Q") action_report ;;
    "-R") shift; action_rollback $1 ;;
    "-U") unmount_all ;;
    *) usage  ;;
esac