#! /bin/bash

# Copyright (C) 2020- The University of Notre Dame
# This software is distributed under the GNU General Public License.
# See the file COPYING for details.

usage() {
    echo "Usage: python_package_run [options] -e <file> command and args ..."
    echo "where options are:"
    echo " -e, --environment <file>   Conda environment as a tar file. (Required.)"
    echo " -d, --unpack-to <dir>      Directory to unpack the environment. If not given,"
    echo "                            a temporary directory is used."
    echo " -w, --wait-for-lock <secs> Number of seconds to wait to get a writing lock"
    echo "                            on <dir>. Default is 300"
    echo " -h, --help                 Show this help screen"
    echo "command and args            Command to execute inside the given environment."
    echo
    exit 1
}

UNPACK_TO=
ENV_NAME=
LOCK_WAIT=300
parse_arguments() {

    original_arg_count=$#

	while [ $# -gt 0 ]
	do
		case $1 in
			-h | --help)
                usage
                ;;
            -d | --unpack-to)
                shift
                UNPACK_TO="$1"
                ;;
            -e | --environment)
                shift
                ENV_NAME="$1"
                ;;
            -w | --wait-for-lock)
                shift
                LOCK_WAIT="$1"
                ;;
            *)
                break
                ;;
        esac
        shift
    done

    if [ -z "${ENV_NAME}" ]
    then
        echo "An environment tarball should be specified with the --environment option."
        usage
        exit 1
    fi

    if [ $# -lt 1 ]
    then
        echo "No command line was specified."
        usage
        exit 1
    fi

    final_arg_count=$#
    arg_counsumed=$((original_arg_count-final_arg_count))

    return ${arg_counsumed}
}

function cleanup {
    if [ "${UNPACK_IS_TMP}" = yes ]
    then
        rm -rf "${UNPACK_TO}"
    fi

    rm -f "${UNTAR_SCRIPT}"
}
trap cleanup EXIT

parse_arguments "$@"


#after shift, whatever is left is taken as the command line to execute.
arg_counsumed=$?
shift ${arg_counsumed}

UNPACK_IS_TMP=no
NEED_TO_UNPACK=yes

if [ -z "${UNPACK_TO}" ]
then
    UNPACK_IS_TMP=yes
    UNPACK_TO="$(mktemp -d)"
fi

UNPACK_TO=$(realpath ${UNPACK_TO})
echo Unpacking conda environment to ${UNPACK_TO}

if [ -f "${UNPACK_TO}" ]
then
    echo "--unpack-to argument is an already existing file. It should be the name of a directory."
    exit 1
fi

UNTAR_SCRIPT=$(mktemp python_package_run.XXXXXX)
cat > ${UNTAR_SCRIPT} << EOF
if [ "\$(find ${UNPACK_TO} -mindepth 1 -maxdepth 1 | wc -l)" = 1 ]
then
    # Directory is empty (the only file there is the lock). It is safe to untar.
    if ! tar -xf "${ENV_NAME}" -C "${UNPACK_TO}"
    then
        echo "Could not uncompress environment: ${ENV_NAME}"
        exit 1
    fi
else
    echo "Directory ${UNPACK_TO} is not empty. Not unpacking environment again."
fi
exit 0
EOF

# unpacking environment if needed.
LOCKFILE="${UNPACK_TO}/.unpacking_python_package_run_lock"
mkdir -p "${UNPACK_TO}"

flock -w ${LOCK_WAIT} ${LOCKFILE} -c "/bin/sh ${UNTAR_SCRIPT}"
if [ "$?" != 0 ]
then
    echo -e "Could not untar environment: ${ENV_NAME}"
    exit 1
fi

#activate and unpack the environment
unset PYTHONPATH
source "${UNPACK_TO}/bin/activate"
if [ "$?" != 0 ]
then
    echo -e "Could not activate environment: ${ENV_NAME}"
    exit 1
fi

flock -w ${LOCK_WAIT} ${LOCKFILE} conda-unpack
if [ "$?" != 0 ]
then
    echo -e "Could not unpack environment: ${ENV_NAME}"
    exit 1
fi

# Finally run the command line:
"${@}"
status=$?

exit $status
