#! /bin/bash

# Copyright (C) 2022 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: poncho_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 [[ "${PONCHO_PACKAGE_RUN_DEBUG}" = yes ]]
    then
        printf "%s poncho_package_run($ORIGINAL_PID): " "$(date +'%Y-%m-%d %T')"
        echo "$@"
    fi
}
export -f logmsg


function errmsg {
    local PONCHO_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 PONCHO_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
        logmsg No command was specified no command will be executed!
        COMMAND_MISSING=yes
    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}")
}


COMMAND_MISSING=no

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

trap cleanup EXIT

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...
if ! mkdir -p "${UNPACK_TO}"
then
	errmsg "could not create directory at ${UNPACK_TO}"
	exit 1
fi



UNTAR_SCRIPT=$(mktemp poncho_package_run.XXXXXX)
cat > "${UNTAR_SCRIPT}" << EOF
if [[ ! -e "${UNPACK_TO}"/.poncho_package_run_expanded ]]
then

    # only untar if it does not look like the directory already has an untared
    # environment. (E.g. another process performed the untar before poncho_package_run)
    if [[ -f ${UNPACK_TO}/bin/conda-unpack ]]
    then
        logmsg directory already expanded but not relocated
    else
        logmsg expanding environment file ${ENV_NAME}
        if ! tar -xf "${ENV_NAME}" -C "${UNPACK_TO}"
        then
            errmsg could not uncompress environment: ${ENV_NAME}
            exit 1
        fi
    fi

    export PATH=${UNPACK_TO}/bin:${PATH}

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

    /bin/date > "${UNPACK_TO}"/.poncho_package_run_expanded
else
    logmsg directory ${UNPACK_TO} is not empty. Not setting 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

# force no args to activate, as it wants to consume "$@" which contains the command line
source "${UNPACK_TO}/bin/activate" ""
if [[ "$?" != 0 ]]
then
    errmsg "could not activate environment: ${ENV_NAME}"
    exit 1
fi

#set enviornment variables
if [[ -e "${UNPACK_TO}/poncho/set_env" ]]
then
    source  "${UNPACK_TO}/poncho/set_env" "${UNPACK_TO}"

fi

# Finally run the command line:
if [[ "${COMMAND_MISSING}" = no ]]
then
    logmsg executing command line: "${@}"
    "${@}"
    status=$?
    logmsg command line exit code: $status
else
    logmsg no command will be executed
fi

exit $status
