#!/usr/bin/env bash
##########################
# This file generates a .tar.bz2 from a .img, as well as compressing the image
# into an .lz4 + .gz file.
# For .bz2 - it will aim to use lbzip2 if available, otherwise fallback to bzip2.
# For .gz - it will aim to use 'pigz' if available, otherwise fallback to gzip
# For .lz4 - it will use lz4 (from liblz4-tool), but will skip if not available
###########################

: ${SOURCE_SELF=0}

DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

: ${FILES_DIR="$DIR"}

cd "$DIR"

: ${IMG_NAME="sg-ubuntu.img"}
: ${OUT_NAME="sg-ubuntu.tar"}

: ${THREADS="$(nproc)"}
: ${MOUNTPT="/mnt/tmp-sg-ubuntu"}

msg() { echo -e "$@"; }
err() { >&2 echo -e "$@"; }

has-cmd() {
    which "$@" &> /dev/null
}

if ! has-cmd pv; then
    err " [!!!] WARNING: Command 'pv' not found. Creating no-op function to passthru pipes to pv..."
    pv() {
        local _zargs _lastarg
        _zargs=("$@")
        if (( $# > 0 )); then
            _lastarg="${_zargs[@]: -1}"
            if ! [[ $_lastarg == ?(-)+([0-9BKGT.]) ]] && [[ -f "$_lastarg" ]]; then
                cat "$f"
                return $?
            fi
        fi
        cat
    }
fi


list-mountpoints() {
    cat /etc/mtab | awk '{ print $2 }'
}

is-mounted() {
    list-mountpoints | grep -Eq "^$MOUNTPT$"
}

has-cmd lbzip2 && export HAS_LBZIP2=1 BZ_CMD="lbzip2" || export HAS_LBZIP2=0 BZ_CMD="bzip2"
has-cmd pigz && export HAS_PIGZ=1 GZ_CMD="pigz" || export HAS_PIGZ=0 GZ_CMD="gzip"

bzip-wrap() {
    _zargs=()
    {
        if (( HAS_LBZIP2 == 0 )); then
            err " [!!!] WARNING: The 'lbzip2' command wasn't found. Falling back to slow single threaded 'bzip2' ..."
            while (( $# > 0 )); do
                if [[ "$1" == "-n" ]]; then
                    shift; shift
                else
                    _zargs+=("$1")
                fi
            done
        else
            _zargs=("$@")
        fi
    } < /dev/null
    "$bz_cmd" "${_zargs[@]}"
}
gzip-wrap() {
    _zargs=()
    {
        if (( HAS_PIGZ == 0 )); then
            err " [!!!] WARNING: The 'pigz' command wasn't found. Falling back to slow single threaded 'bzip2' ..."
            while (( $# > 0 )); do
                if [[ "$1" == "-p" || "$1" == "--processes" ]]; then
                    shift; shift
                else
                    _zargs+=("$1")
                fi
            done
        else
            _zargs=("$@")
        fi
    } < /dev/null
    "$gz_cmd" "${_zargs[@]}"
}

do-umount() {
    if is-mounted "$1"; then
        >&2 umount -v "$@"
        _ret=$?
        if (( _ret )); then
            err " [!!!] Non-zero return code from 'umount -v $@' - code: $_ret"
            err " [!!!] Aborting unmounting...\n"
            return $_ret
        fi
    fi
    return 0
}

cmp-tar() {
    if ! is-mounted "$MOUNTPT"; then
        >&2 mount -v -o ro "$IMG_NAME" "$MOUNTPT"
        _ret=$?
        if (( _ret )); then
            err " [!!!] Non-zero return code from 'mount -v -o ro $IMG_NAME $MOUNTPT' - code: $_ret"
            err " [!!!] Aborting tarballing...\n"
            return $_ret
        fi
    fi
    cd "$MOUNTPT"
    _img_size="$(stat -c%s "$IMG_NAME")"
    _out_name="${OUT_NAME}.bz2"
    msg " >>> Tarring current dir ${PWD} and piping it into 'bzip-wrap -9 -n $THREADS'"
    msg "   > Output file: $_out_name"
    tar --one-file-system -cpf - . | pv -s "$_img_size" | bzip-wrap -9 -n $THREADS > "$_out_name"
    _ret=$?
    cd - &> /dev/null
    do-umount "$MOUNTPT"
    if (( _ret )); then
        err " [!!!] Non-zero return code from 'tar --one-file-system -cpf - . | bzip-wrap -9 -n $THREADS'"
        err " [!!!] Code: $_ret || Aborting tarballing...\n"
        return $_ret
    fi

    msg " +++ Finished tarring current dir ${PWD} and piping it into 'bzip-wrap -9 -n $THREADS'"
    return 0
}


cmp-img-gz() {
    cd "$FILES_DIR"
    _img_size="$(stat -c%s "$IMG_NAME")"
    _out_name="${IMG_NAME}.gz"
    msg " >>> Compressing image ${PWD}/${IMG_NAME} with gzip."
    msg "   > Output file: $_out_name"
    pv -s "$_img_size" < "$IMG_NAME" | gzip-wrap -1 -v -p $THREADS > "$_out_name"
    _ret=$?
    if (( _ret )); then
        err " [!!!] Non-zero return code from 'gzip-wrap -1 -v -p $THREADS > \"$_out_name\"'"
        err " [!!!] Code: $_ret || Aborting compression...\n"
        return $_ret
    fi

    msg " +++ Finished compressing '$IMG_NAME' with gzip into $_out_name"
    return 0
}

cmp-img-lz4() {
    if ! has-cmd lz4; then
        err " [!!!] WARNING: Command 'lz4' not found. Skipping compressing image with lz4..."
        return 1
    fi

    cd "$FILES_DIR"
    _img_size="$(stat -c%s "$IMG_NAME")"
    _out_name="${IMG_NAME}.lz4"
    msg " >>> Compressing image ${PWD}/${IMG_NAME} with lz4."
    msg "   > Output file: $_out_name"
    lz4 -k -v --sparse -z -9 "$IMG_NAME" -c > "$_out_name" 
    _ret=$?
    if (( _ret )); then
        err " [!!!] Non-zero return code from 'lz4 -k -v --sparse -z -9 \"$IMG_NAME\" -c > \"$_out_name\"'"
        err " [!!!] Code: $_ret || Aborting compression...\n"
        return $_ret
    fi

    msg " +++ Finished compressing '$IMG_NAME' with lz4 into $_out_name"
    return 0
}

script-main() {
    cd "$FILES_DIR"
    if ! [[ -f "$IMG_NAME" ]]; then
        err " [!!!] The image file '${IMG_NAME}' doesn't exit - so we have nothing to base the archives on."
        err " [!!!] Cannot continue. Please ensure that '${FILES_DIR}/${IMG_NAME}' exists before calling this script."
        return 4
    fi

    if ! [[ -f "${OUT_NAME}.bz2" ]]; then
        msg " >>> Tarball ${OUT_NAME}.bz2 doesn't exist - calling cmp-tar to create it!"
        cmp-tar
    fi
    do-umount "$MOUNTPT"
    if ! [[ -f "${IMG_NAME}.lz4" ]]; then
        msg " >>> Compressed image ${IMG_NAME}.lz4 doesn't exist - calling cmp-img-lz4 to create it!"
        cmp-img-lz4
    fi
    if ! [[ -f "${IMG_NAME}.gz" ]]; then
        msg " >>> Compressed image ${IMG_NAME}.gz doesn't exist - calling cmp-img-gz to create it!"
        cmp-img-lz4
    fi
}

if (( SOURCE_SELF )); then
    err " +++ SOURCE_SELF is true. Not calling script-main"
else
    script-main "$@"
fi

