#!/bin/bash

# Wrap zpaq as a more generic compression utility
#
# Copyright 2017 Robert Krawitz (rlk@alum.mit.edu)
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

declare -r ZPAQ=/usr/bin/zpaq
declare METHOD=5
declare VERBOSE=
declare ZPCAT=
declare DECOMPRESS=
declare FORCE=
declare KEEP=
declare SAVE_META=
declare RECURSE=
declare STATUS=0
declare -r PROGNAME=$0
declare -r BASEDIR=$(pwd)
declare TMP_FILE=
declare TMP_DIR=
declare TMP_ARCHIVE=

usage() {
    declare -l USTATUS=${1:-0}
    echo "Usage: $PROGNAME [OPTIONS] [FILE]..."
    echo "Compress or decompress FILEs (by default, compress FILEs in place)"
    echo
    echo "  -c   Write to standard output, keep original files unchanged"
    echo "  -d   Decompress"
    echo "  -f   Force overwrite of output file"
    echo "  -h   Print this help"
    echo "  -k   Keep (don't delete) input files"
    echo "  -n	 Do not preserve original filename and timestamp"
    echo "  -N	 Preserve original filename and timestamp"
    echo "  -q   Suppress all messages"
    echo "  -r   Operate recursively on directories"
    echo "  -v   Print verbose messages"
    echo "  -1   Use minimum compression"
    echo "  -5   Use maximum compression"
    echo
    echo "With no FILE, or when only FILE is -, read standard input"
    exit "$USTATUS"
}

while getopts "qvcdhfknNrz0123456789" opt ; do
    case "$opt" in
	q)     VERBOSE=			;;
	v)     VERBOSE=1		;;
	c)     ZPCAT=1			;;
	d)     DECOMPRESS=1		;;
	f)     FORCE=1			;;
	h)     usage			;;
	k)     KEEP=1			;;
	n)     SAVE_META=-noattributes	;;
	N)     SAVE_META=		;;
	r)     RECURSE=1		;;
	z)     DECOMPRESS=		;;
	[0-9]) METHOD=$opt	 	;;
	*)     usage 1			;;
    esac
done
shift $((OPTIND-1))

if [[ -n $RECURSE && -n $ZPCAT ]] ; then
    echo "May not combine recursive and output to stdout"
    usage "$0" 1
fi

if [[ -n $RECURSE && -z $* ]] ; then
    echo "May not use recursive with no inputs"
    usage "$0" 1
fi

declare -a FILES

g() {
    if [[ $1 == /* ]] ; then
	echo "$1"
    else
	echo "$BASEDIR/$1"
    fi
}

build_file_list() {
    for f in "$@" ; do
	if [[ -d $f ]] ; then
	    if [[ -n $RECURSE ]] ; then
		OIFS=$IFS
		IFS=
		if [[ -n $DECOMPRESS ]] ; then
		    while read -r FILE ; do
			FILES+=("$(g "$f1")")
		    done <<< "$(find "$f" -type f -name '*.zpaq' -print)"
		else
		    while read -r FILE ; do
			FILES+=("$(g "$f1")")
		    done <<< "$(find "$f" -type f \! -name '*.zpaq' -print)"
		fi
		IFS=$OIFS
	    else
		echo "$f: is a directory, skipping" 1>&2
		STATUS=1
	    fi
	elif [[ -f $f ]] ; then
	    if [[ -n $DECOMPRESS && $f != *.zpaq ]] ; then
		echo "$f: not a zpaq archive, skipping" 1>&2
		STATUS=1
	    elif [[ -z $DECOMPRESS && $f == *.zpaq ]] ; then
		echo "$f: is a zpaq archive, skipping" 1>&2
		STATUS=1
	    else
		FILES+=("$(g "$f")")
	    fi
	else
	    echo "$f: not a plain file, skipping" 1>&2
	    STATUS=1
	fi
    done
}

if [[ -z $* || $* == '-' ]] ; then
    ZPCAT=1
else
    build_file_list "$@"
fi

if [[ -n $ZPCAT && ${#FILES[@]} -gt 1 ]] ; then
    echo "May not compress/decompress more than one file to stdout" 1>&2
    usage 1
fi

# Wrappers to simplify debugging

run_zpaq() {
    if [[ -n $VERBOSE ]] ; then
	"$ZPAQ" "$@"
    else
	"$ZPAQ" "$@" >/dev/null 2>&1
    fi
}

RM() {
    for f in "$@" ; do
	if [[ $1 != "$TMP_FILE" && $1 != "$TMP_DIR" && $1 != "$TMP_ARCHIVE" ]] ; then
	    rm -rf "$f"
	fi
    done
}

do_decompress() {
    FILE="$1"
    DEST="$2"
    # shellcheck disable=SC2164
    [[ -n $3 && -d $3 ]] && cd "$3"
    run_zpaq extract "$FILE" $SAVE_META || STATUS=1
    if [[ $(wc -l <<< "$(find . -type f -print)") -gt 1 ]] ; then
	echo "Multiple file archive $TMP_ARCHIVE, skipping" 1>&2
	exit 1
    fi
    if [[ -z $DEST ]] ; then
	cat "$(find . -type f -print)"
    else
	mv "$(find . -type f -print)" "$DEST"
    fi
}

decompress_file() {
    FILE="$1"
    DEST="$2"
    if [[ -d $TMP_DIR ]] ; then
	# Decompress into a temporary directory.  We don't actually
	# know what's in the archive and don't want to dump arbitrary
	# contents into the user's directory.  We do this in a subshell
	# so that we don't have to try to cd back to the current directory
	# (which may be problematic if the directory is remote and we
	# spend a lot of time in the temp directory).
	(do_decompress "$FILE" "$DEST" "$TMP_DIR")
	# shellcheck disable=2181
	if (( $? > 0 )) ; then
	    STATUS=1
	elif [[ -z $ZPCAT && -z $KEEP ]] ; then
	    RM "$FILE"
	fi
    fi
}

decompress() {
    # Decompress: decompress all .zpaq files.
    if [[ -z ${FILES[*]} ]] ; then
	# Decompress stdin implies output to stdout
	ZPCAT=1
	cat > "$TMP_ARCHIVE" && decompress_file "$TMP_ARCHIVE"
    else
	for f in "${FILES[@]}" ; do
	    declare DEST=
	    [[ -z $ZPCAT ]] && DEST="${f%.zpaq}"
	    if [[ -n $DEST && -e $DEST && -z $FORCE ]] ; then
		echo "${DEST#${BASEDIR}/} exists, skipping" 1>&2
		STATUS=1
	    else
		decompress_file "$f" "$DEST"
	    fi
	done
    fi
}

do_compress() {
    FILE="$1"
    DEST=$(g "$2")
    cd "${FILE%/*}" || return 1
    run_zpaq add "$DEST" "${FILE##*/}" $SAVE_META -method "$METHOD"
}

compress_file() {
    FILE="$1"
    DEST=${2:-${FILE}.zpaq}
    # Make sure that DEST really is empty.
    RM "$DEST"
    (do_compress "$FILE" "$DEST") || return 1
    if [[ -n "$ZPCAT" ]] ; then
	# zpaq won't send anything to stdout; we have to
	cat "$DEST"
	RM "$DEST"
    elif [[ -z $KEEP ]] ; then
	RM "$FILE"
    fi
    return 0
}

compress() {
    if [[ -z ${FILES[*]} ]] ; then
	# Compress stdin implies output to stdout
	ZPCAT=1
	cat > "$TMP_FILE" && {
	    compress_file "$TMP_FILE" "$TMP_ARCHIVE" || STATUS=1
	}
    else
	for f in "${FILES[@]}" ; do
	    declare DEST="${f}.zpaq"
	    [[ -n $ZPCAT ]] && DEST="$TMP_ARCHIVE"
	    if [[ -z $ZPCAT && -e $DEST && -z $FORCE ]] ; then
		echo "${DEST#${BASEDIR}/} exists, skipping" 1>&2
		STATUS=1
	    else
		compress_file "$f" "$DEST" || STATUS=1
	    fi
	done
    fi
}

# If all of the command line files were disqualified because they
# are compressed files, don't try to compress stdin.
[[ -n $* && -z ${FILES[*]} ]] && exit 1

cleanup() {
    [[ -n $TMP_FILE && -f $TMP_FILE ]] && rm -f "$TMP_FILE"
    [[ -n $TMP_ARCHIVE && -f $TMP_ARCHIVE ]] && rm -f "$TMP_ARCHIVE"
    [[ -n $TMP_DIR && -d $TMP_DIR ]] && rm -rf "$TMP_DIR"
}

trap cleanup EXIT

declare TMP_FILE=$(mktemp /tmp/wzpaqinXXXXXXXXX)
declare TMP_DIR=$(mktemp -d /tmp/wzpaq_outXXXXXXXXX)
declare TMP_ARCHIVE=$(mktemp /tmp/wzpaqXXXXXXXXX.zpaq)

if [[ -n $DECOMPRESS ]] ; then
    decompress
else
    compress
fi

exit $STATUS