#! /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.

# Executes a command line inside a conda environment.
# The conda environment can be supplied via a tar file or a directory.
#
# It is assumed that the conda environment can be used in the absence of a
# local conda installation. This means it should include conda and conda-unpack
# itself. This requirement can be fulfilled with the generic recipe shown with
# the --help-env-creation option:
#

start_time=$(date +'%s%N')

usage() {
    echo "Usage: python_package_run [options] -e <file> command and args ..."
    echo "where options are:"
    echo " -e, --environment <path>   Conda environment as a tar file or a directory. (Required.)"
    echo " -u, --unpack-to <dir>      Directory to unpack the environment. If not given,"
    echo "                            a temporary directory is used. If the argument to"
    echo "                            --environment is a directory, --unpack-to is ignored."
    echo " -w, --wait-for-lock <secs> Number of seconds to wait to get a writing lock"
    echo "                            on <dir>. Default is 300"
    echo " -d, --debug                Print debug messages."
    echo " --help-env-creation        Show instructions to create conda environments as tar files."
    echo " -h, --help                 Show this help screen."
    echo "command and args            Command to execute inside the given environment."
    echo
}

function help_env_creation {
cat << EOF

Environments should include conda and conda-unpack so they can be used
in the absence of a local conda installation. Environment creation follows
the generic recipe:

$ conda create --prefix ./my-conda-env python=X.XX conda
$ source ./my-conda-env/bin/activate
$ conda install -c conda-forge conda-pack
... conda and pip install rest of packages ...
$ python -c 'import conda_pack; conda_pack.pack(prefix="my-conda-env")

The above generates my-conda-env.tar.gz ready to be used by this script as:

$0 --environment my-conda-env.tar.gz -- python -c 'print(42)'

The directory ./my-conda-env can be safely removed.

EOF
}

export ORIGINAL_PID=$$

function logmsg {
    if [[ "${PYTHON_PACKAGE_RUN_DEBUG}" = yes ]]
    then
        printf "%s python_package_run($ORIGINAL_PID): " "$(date +'%Y-%m-%d %T')"
        echo "$@"
    fi
}
export -f logmsg


function errmsg {
    local PYTHON_PACKAGE_RUN_DEBUG=yes
    logmsg "$@"
}
export -f errmsg


UNPACK_TO=
ENV_NAME=
LOCK_WAIT=300
parse_arguments() {

    original_arg_count=$#

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

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

    if [[ $# -lt 1 ]]
    then
        errmsg "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}"

    end_time=$(date +'%s%N')

    total_ns=$((end_time-start_time))
    total_ms=$((total_ns/1000000))
    ms=$((total_ms % 1000))
    s=$((total_ms/1000))

    logmsg total runtime: $(printf "%d.%0.3d seconds" ${s} ${ms})
}
trap cleanup EXIT

parse_arguments "$@"

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

ENV_NAME_IS_DIR=no
UNPACK_IS_TMP=no
NEED_TO_UNPACK=yes

if [[ -d "${ENV_NAME}" ]]
then
    ENV_NAME_IS_DIR=yes

    if [[ -n "${UNPACK_TO}" ]]
    then
        errmsg "ignoring --unpack-to argument as --environment is a directory." 
    fi

    UNPACK_TO="${ENV_NAME}"
fi

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

UNPACK_TO=$(realpath ${UNPACK_TO})
logmsg setting up conda environment at ${UNPACK_TO}

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

# unpacking environment if needed...
mkdir -p "${UNPACK_TO}"

UNTAR_SCRIPT=$(mktemp python_package_run.XXXXXX)
cat > ${UNTAR_SCRIPT} << EOF
if [[ ! -e "${UNPACK_TO}"/.python_package_run_expanded ]]
then
    logmsg expanding environment file ${ENV_NAME}

    # Directory is empty. It is safe to untar.
    if ! tar -xf "${ENV_NAME}" -C "${UNPACK_TO}"
    then
        errmsg could not uncompress environment: ${ENV_NAME}
        exit 1
    fi

    if ! ${UNPACK_TO}/bin/conda-unpack
    then
        errmsg could not relocate environment: ${ENV_NAME}
        exit 1
    fi

    /bin/date > "${UNPACK_TO}"/.python_package_run_expanded
    /bin/sync "${UNPACK_TO}"
else
    logmsg directory ${UNPACK_TO} is not empty. Not expanding environment file again
fi
exit 0
EOF

if [[ "${ENV_NAME_IS_DIR}" = no ]]
then
    logmsg waiting for lock on directory ${UNPACK_TO}
    flock -E 222 -w ${LOCK_WAIT} ${UNPACK_TO} -c "/bin/bash ${UNTAR_SCRIPT}"
    status=$?
    if [[ "$status" != 0 ]]
    then
        if [[ "$status" = 222 ]]
        then
            errmsg could not acquire lock
        else
            errmsg could not untar environment: ${ENV_NAME}
        fi
        exit 1
    fi
fi

#activate and unpack the environment
logmsg activating environment
unset PYTHONPATH
source "${UNPACK_TO}/bin/activate"
if [[ "$?" != 0 ]]
then
    errmsg "could not activate environment: ${ENV_NAME}"
    exit 1
fi

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

logmsg command line exit code: $status

exit $status
