#!/bin/bash

# copy stdout and stderr to "$logfile"
logfile='./create_env.log'
exec > >(tee -i $logfile)
exec 2>&1

# move the logfile to the `prefix.log`, if set
function atexit {
    if test -z "$PREFIX"; then
        mv $logfile $logfile.err
    else
        echo "move logfile to 'env_$(basename $PREFIX).log'"
        mv $logfile "env_$(basename $PREFIX).log"
    fi
}
trap atexit EXIT

echo "[$0 $*]"

ru_env=$(which radical-utils-env.sh)
if test -z '$ru_env'
then
    echo "$0 needs radical.utils installed"
    exit 1
fi

. radical-utils-env.sh

set -e  # exit on all errors

SCRIPT="$0"
PREFIX=       # full path or virtenv name
VERSION=      # python version (e.g., 3.8, 3.11 etc.)
              # NOTE: for virtualenv corresponding version should be installed
MODULES=      # alternative list of pre-installed modules/packages
DEFAULTS=     # install default modules
ETYPE='venv'  # type of environment: 'virtualenv'/'venv' or 'anaconda'/'conda'

pre_exec(){
    echo "pre exec: $*"
    eval $*
    echo "pre exec done"
}

while getopts "e:p:P:v:m:t:T:d" OPTION; do
    case $OPTION in
        p)  PREFIX="$OPTARG"    ;;
        P)  pre_exec "$OPTARG"  ;;
        v)  VERSION="$OPTARG"   ;;
        m)  MODULES="$OPTARG"   ;;
        t)  ETYPE="$OPTARG"     ;;
        T)  EDUMP="$OPTARG"     ;;
        d)  DEFAULTS=True       ;;
        *)  echo "Unknown option: '$OPTION'='$OPTARG'"
            exit 1;;
    esac
done

test -z "$PREFIX"  && echo "no prefix" && exit

test "$ETYPE" = 'anaconda'   && ETYPE='conda'
test "$ETYPE" = 'virtenv'    && ETYPE='venv'
test "$ETYPE" = 'virtualenv' && ETYPE='venv'

# sanitize name (remove trailing slashes etc)
PREFIX="$(dirname $PREFIX)/$(basename $PREFIX)"


# move from comma separated to space separated module list
MODULES=$(echo "$MODULES" | sed -e 's/,/ /g')

if ! test -z "$DEFAULTS"
then
    # by default, install all RCT dependencies
    MODULES="$MODULES chardet colorama idna msgpack"
    MODULES="$MODULES msgpack-python netifaces ntplib parse dill"
    MODULES="$MODULES pyzmq regex requests setproctitle urllib3"
fi


# ------------------------------------------------------------------------------
#
help(){

    err="$1"
    ret=0

    if ! test -z "$err"
    then
        ret=1
        printf "\n    ERROR: $err\n"
    fi

    cat <<EOT

    $0

    This script creates a virtual environment (conda or venv) at the given
    target path.  That environment should be suitable to be used as static env
    for a radical.pilot target resource, and can be specified in a resource
    configuration for RP.

    Options:
      -p <prefix>    - path to use for environment
      -P <pre_exec>  - command to be run before creation
      -v <version>   - python version to be used
      -m <modules>   - comma separated list of modules to install
      -t <type>      - environment type to create (venv, conda)
      -T <env_path>  - target file where to dump the environment
      -d             - install default RCT dependencies into the new environment

EOT
    exit $ret
}


# ------------------------------------------------------------------------------
#
test    "$PREFIX" = "-h" && help
test -z "$PREFIX"        && help "missing target"

# Ensure we install in an absolute path -- pip seems to like that better...
case $PREFIX in
    /*)
        ;;
    *)
        PREFIX="$(pwd)/$PREFIX"
        ;;
esac

PYTHON=python3

if test -z "$VERSION"
then
    VERSION=$($PYTHON -V | cut -f 2 -d ' ' | cut -f 1,2 -d '.')
fi

# we only ever use major and minor version numbers, and will ignore subminor
VERSION=$(echo $VERSION | cut -f 1,2 -d '.')

# check env type
if test "$ETYPE" = 'shell'
then
    # nothing to do for a plain shell environment
    true

elif test "$ETYPE" = 'conda'
then

    if test -d "$PREFIX"
    then

        printf "env at $PREFIX exists\n"
        if ! test -z "$(which conda)"; then
            eval "$(conda shell.posix hook)"
            conda activate "$PREFIX"  # ?! test

        elif test -e "$PREFIX/bin/activate"; then
            . "$PREFIX/bin/activate"

        elif test -e "$PREFIX/../../bin/activate"; then
            # check conda base directory
            . "$PREFIX/../../bin/activate" "$(basename $PREFIX)"

        else
            help "conda not found"
        fi

    else

        which conda || help "conda not found"
        eval "$(conda shell.bash hook)"
        conda activate base

        printf "create  $ETYPE ve [$PREFIX]\n"
        if test -z "$VERSION"
        then
            conda create -y -p "$PREFIX"
        else
            conda create -y -p "$PREFIX" "python=$VERSION"
        fi
        conda activate "$PREFIX"

        echo -n "update  setuptools pip wheel"
        conda install setuptools pip wheel || exit 1

        for mod in $MODULES
        do
            # use `pip install` instead of `conda install` if `mod` is a path
            case "$mod" in
                /*|*/*)
                    echo -n "install $mod with pip"
                    pip install "$mod" || exit 1
                    ;;
                *)
                    echo -n "install $mod with conda"
                    conda install "$mod" \
                    || pip install "$mod" || exit 1
                    ;;
            esac
        done
    fi

elif test "$ETYPE" = 'venv'
then
    if ! test -z "$VERSION"
    then
        PY_VERSION=$($PYTHON -V | cut -f 2 -d ' ' | cut -f 1,2 -d '.')
        if ! test "$VERSION" = "$PY_VERSION"
        then
            echo "requested version $VERSION not supported for venv ($PY_VERSION)"
            exit 1
        fi
    fi

    if test -d "$PREFIX"
    then
        printf "\nenv at $PREFIX exists\n"
        . "$PREFIX"/bin/activate
    else

        # TODO: find out why this hangs in GH ci
        # VIRTENV_CMD="$(which virtualenv 2>/dev/null)"
        VIRTENV_CMD="$PYTHON -m venv"

        if test -z "$VIRTENV_CMD"
        then
            echo -n "install private virtualenv "
            VIRTENV_VER=20.0.27
            VIRTENV_TGZ="$VIRTENV_VER.tar.gz"
            VIRTENV_TGZ_URL="https://github.com/pypa/virtualenv/archive/$VIRTENV_TGZ"

            curl -k -L -O "$VIRTENV_TGZ_URL"
            tar zxf "$VIRTENV_TGZ"
            VIRTENV_CMD="$PYTHON $VIRTENV_VER/virtualenv.py"
        fi

        echo -n "create  virtualenv "
        $VIRTENV_CMD "$PREFIX"

        . "$PREFIX"/bin/activate

        # ensure we have pip
        python3 -m ensurepip

        echo -n "update  setuptools pip wheel"
        pip install --no-cache-dir --upgrade setuptools pip wheel || exit 1

        for mod in $MODULES
        do
            echo -n "install $mod with pip"
            pip install --no-cache-dir "$mod" || exit 1
        done
    fi
fi

if test -z "$EDUMP"
then
    mkdir -p "$(pwd)/env"
    EDUMP="$(pwd)/env/$(basename $PREFIX).env"
fi

env_dump -t "$EDUMP"

# print the ve information and stack versions for verification
echo
echo "---------------------------------------------------------------------"
echo
echo "PWD       : $(pwd)"
echo "ENV       : $EDUMP"
echo "TYPE      : $ETYPE"
echo "SCRIPT    : $SCRIPT"
echo "PREFIX    : $PREFIX"
echo "VERSION   : $VERSION"
echo "MODULES   : $MODULES"
echo "DEFAULTS  : $DEFAULTS"
echo "PYTHON    : $(which $PYTHON) ($($PYTHON -V))"
echo "PYTHONPATH: $PYTHONPATH"
echo "RCT_STACK : $(radical-stack || true)"
echo
echo "---------------------------------------------------------------------"
echo

