#!/usr/bin/env sh
# NOTE: This script should be sourced! The shebang is only here to help syntax highlighters.

if ! (return 0 2> /dev/null); then
    echo "ERROR: Source this script: source '$0'." >&2
    exit 1
fi

cleanup() {
    unset _ARCH
    unset _BASEEXE
    unset _CHANNEL_NAME
    unset _CMD
    unset _DEFAULT_ARCH
    unset _DEFAULT_DEVENV
    unset _DEFAULT_DRYRUN
    unset _DEFAULT_INSTALLER
    unset _DEFAULT_MACH
    unset _DEFAULT_PYTHON
    unset _DEFAULT_UPDATE
    unset _DEVENV
    unset _DOWNLOAD_URL
    unset _DRYRUN
    unset _ENV
    unset _ENVEXE
    unset _INSTALLER
    unset _INSTALLER_FILE
    unset _INSTALLER_INITIAL
    unset _INSTALLER_TYPE
    unset _MACH
    unset _NAME
    unset _PYTHON
    unset _PYTHONEXE
    unset _SRC
    unset _UPDATE
    unset _UPDATED
}
updating() {
    # check if explicitly updating or if 24 hrs since last update
    [ ${_UPDATE} = "true" ] && return 0
    [ -f "${_UPDATED}" ] || return 0
    return $(( $(( $(date +%s) - $(date -r "${_UPDATED}" +%s) )) < 86400 ))
}

cleanup

# get source path
# since zsh is MacOS standard fallback to $0 if not in bash
_SRC="$(cd "$(dirname "$(dirname "${BASH_SOURCE:-$0}")")" 2>&1 > /dev/null; pwd -P)"

# define default values
_DEFAULT_PYTHON="3.10"
_DEFAULT_INSTALLER="miniconda"
# $ uname -sm
# Linux x86_64
# Linux aarch64
# Darwin arm64
# ...
_DEFAULT_MACH="$(uname)"
_DEFAULT_ARCH="$(uname -m)"
_DEFAULT_UPDATE="false"
_DEFAULT_DEVENV="${_SRC}/devenv"
_DEFAULT_DRYRUN="false"

# parse args
while [ $# -gt 0 ]; do
    case $1 in
        -p|--python)
            _PYTHON="$2"
            shift
            shift
            ;;
        --python=*)
            _PYTHON="${1#*=}"
            shift
            ;;
        -m|--mach)
            _MACH="$2"
            shift
            shift
            ;;
        --mach=*)
            _MACH="${1#*=}"
            shift
            ;;
        -a|--arch)
            _ARCH="$2"
            shift
            shift
            ;;
        --arch=*)
            _ARCH="${1#*=}"
            shift
            ;;
        -i|--installer)
            _INSTALLER_TYPE="$2"
            shift
            shift
            ;;
        --installer=*)
            _INSTALLER_TYPE="${1#*=}"
            shift
            ;;
        -u|--update)
            _UPDATE="true"
            shift
            ;;
        -d|--devenv)
            _DEVENV="$2"
            shift
            shift
            ;;
        --devenv=*)
            _DEVENV="${1#*=}"
            shift
            ;;
        -n|--dry-run)
            _DRYRUN="true"
            shift
            ;;
        -h|--help)
            # since zsh is MacOS standard fallback to $0 if not in bash
            echo "Usage: source ${BASH_SOURCE:-$0} [options]"
            echo ""
            echo "Options:"
            echo "  -p, --python    VERSION    Python version for the env to activate. (default: ${_DEFAULT_PYTHON})"
            echo "  -i, --installer INSTALLER  Installer to use: miniconda or miniforge, can also be defined in ~/.condarc. (default: ${_DEFAULT_INSTALLER})"
            echo "  -m, --mach      MACHINE    Machine type for the env to activate. (default: ${_DEFAULT_MACH})"
            echo "  -a, --arch      ARCH       Architecture for the env to activate. (default: ${_DEFAULT_ARCH})"
            echo "  -u, --update               Force update packages. (default: ${_DEFAULT_UPDATE}, update every 24 hours)"
            echo "  -d, --devenv    PATH       Path to base env install, can also be defined in ~/.condarc."
            echo "                             Path is appended with machine hardware name, see --mach. (default: ${_DEFAULT_DEVENV})"
            echo "  -n, --dry-run              Display env to activate. (default: ${_DEFAULT_DRYRUN})"
            echo "  -h, --help                 Display this."
            return 0
            ;;
        *)
            echo "Error: unknown option $1" >&2
            return 1
            ;;
    esac
done

# read values key from ~/.condarc if not set
if [ -f ~/.condarc ]; then
    _INSTALLER_TYPE="${_INSTALLER_TYPE:-$(grep "^installer_type:" ~/.condarc | tail -n 1 | sed -e "s/^.*:[[:space:]]*//" -e "s/[[:space:]]*$//")}"
    _DEVENV="${_DEVENV:-$(grep "^devenv:" ~/.condarc | tail -n 1 | sed -e "s/^.*:[[:space:]]*//" -e "s/[[:space:]]*$//")}"
fi

# prompt for installer type if not set
if [ -z "${_INSTALLER_TYPE}" ]; then
    echo "Choose conda installer:"
    echo "  1) miniconda (default - Anaconda defaults channel)"
    echo "  2) miniforge (conda-forge channel)"
    echo ""
    echo "Note: This choice can be overridden by setting the 'installer_type' key in ~/.condarc."
    echo ""
    printf "Enter choice [1]: "
    read -r _INSTALLER_TYPE
    # normalize user input
    case "${_INSTALLER_TYPE}" in
        1|"") _INSTALLER_TYPE="miniconda" ;;
        2) _INSTALLER_TYPE="miniforge" ;;
        *)
            echo "Error: invalid choice '${_INSTALLER_TYPE}'. Please run again and choose 1 or 2." >&2
            return 1
    esac
fi

# fallback to default values if not set
_PYTHON="${_PYTHON:-3.10}"
_INSTALLER_TYPE="${_INSTALLER_TYPE:-${_DEFAULT_INSTALLER}}"
_DEVENV="${_DEVENV:-${_DEFAULT_DEVENV}}"
# $ uname -sm
# Linux x86_64
# Linux aarch64
# Darwin arm64
# ...
_MACH="${_MACH:-${_DEFAULT_MACH}}"
_ARCH="${_ARCH:-${_DEFAULT_ARCH}}"
_UPDATE="${_UPDATE:-${_DEFAULT_UPDATE}}"
_DRYRUN="${_DRYRUN:-${_DEFAULT_DRYRUN}}"

# validate installer type
case "${_INSTALLER_TYPE}" in
    miniconda|miniforge) ;;
    *)
        echo "Error: invalid installer type '${_INSTALLER_TYPE}'. Must be 'miniconda' or 'miniforge'." >&2
        return 1
        ;;
esac

# devenv tilde expansion
_DEVENV="${_DEVENV/#\~/${HOME}}"

# installer location is root devenv directory
_INSTALLER="${_DEVENV}/installers/${_MACH}/${_ARCH}"

# devenv include OS
_DEVENV="${_DEVENV}/${_MACH}/${_ARCH}"

# other environment values
_NAME="devenv-${_PYTHON}-${_INSTALLER_TYPE}"
_ENV="${_DEVENV}/envs/${_NAME}"
_UPDATED="${_ENV}/.devenv-updated"
case "${_MACH}" in
    Darwin|Linux) _BASEEXE="${_DEVENV}/bin/conda" ; _ENVEXE="${_ENV}/bin/conda" ; _PYTHONEXE="${_ENV}/bin/python" ;;
               *) _BASEEXE="${_DEVENV}/Scripts/conda.exe" ; _ENVEXE="${_ENV}/Scripts/conda.exe" ; _PYTHONEXE="${_ENV}/Scripts/python.exe" ;;
esac

# set installer-specific values
case "${_INSTALLER_TYPE}" in
    miniconda)
        _INSTALLER_FILE="${_INSTALLER}/miniconda"
        case "${_MACH}" in
            Darwin) _DOWNLOAD_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-${_ARCH}.sh" ;;
            Linux)  _DOWNLOAD_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${_ARCH}.sh" ;;
            *)      _DOWNLOAD_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe" ;;
        esac
        ;;
    miniforge)
        _INSTALLER_FILE="${_INSTALLER}/miniforge"
        case "${_MACH}" in
            Darwin) _DOWNLOAD_URL="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-MacOSX-${_ARCH}.sh" ;;
            Linux)  _DOWNLOAD_URL="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-${_ARCH}.sh" ;;
            *)      _DOWNLOAD_URL="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Windows-x86_64.exe" ;;
        esac
        ;;
esac
case "${_MACH}" in
    Darwin|Linux) _INSTALLER_FILE="${_INSTALLER_FILE}.sh" ;;
    *)            _INSTALLER_FILE="${_INSTALLER_FILE}.exe" ;;
esac

# dryrun printout
if [ ${_DRYRUN} = "true" ]; then
    echo "Python: ${_PYTHON}"
    echo "Machine: ${_MACH}"
    echo "Architecture: ${_ARCH}"
    echo "Installer [${_INSTALLER_TYPE}]: ${_INSTALLER_FILE} $([ -f "${_INSTALLER_FILE}" ] && echo "[exists]" || echo "[pending]")"
    echo "Devenv: ${_DEVENV} $([ -e "${_DEVENV}" ] && echo "[exists]" || echo "[pending]")"
    echo "Updating: $(updating && echo "[yes]" || echo "[no]")"
    echo ""
    echo "Name: ${_NAME}"
    echo "Path: ${_ENV} $([ -e "${_ENV}" ] && echo "[exists]" || echo "[pending]")"
    echo ""
    echo "Source: ${_SRC}"
    return 0
fi

# ensure installer and devenv directories exist
mkdir -p "${_INSTALLER}"
mkdir -p "${_DEVENV}"
_INSTALLER="$(cd "${_INSTALLER}" 2>&1 > /dev/null ; pwd -P)"
_DEVENV="$(cd "${_DEVENV}" 2>&1 > /dev/null ; pwd -P)"

# deactivate any prior envs
if ! [ ${CONDA_SHLVL:-0} = 0 ]; then
    echo "Deactivating ${CONDA_SHLVL} environment(s)..."
    while ! [ ${CONDA_SHLVL:-0} = 0 ]; do
        if ! conda deactivate; then
            echo "Error: failed to deactivate environment(s)" 1>&2
            return 1
        fi
    done
fi

# does conda install exist?
if ! [ -f "${_DEVENV}/conda-meta/history" ]; then
    # Remove zero-byte installer files before download
    if [ -f "${_INSTALLER_FILE}" ] && [ ! -s "${_INSTALLER_FILE}" ]; then
        echo "Warning: removing empty installer file ${_INSTALLER_FILE}"
        rm -f "${_INSTALLER_FILE}"
    fi

    # downloading installer
    if ! [ -f "${_INSTALLER_FILE}" ]; then
        echo "Downloading ${_INSTALLER_TYPE}..."
        if [ "${_MACH}" = "Darwin" ] || [ "${_MACH}" = "Linux" ]; then
            curl -L --fail --silent --show-error "${_DOWNLOAD_URL}" -o "${_INSTALLER_FILE}"
        else
            powershell.exe -Command "Invoke-WebRequest -Uri '${_DOWNLOAD_URL}' -OutFile '${_INSTALLER_FILE}' | Out-Null"
        fi
        if ! [ $? = 0 ] || [ ! -s "${_INSTALLER_FILE}" ]; then
            echo "Error: failed to download ${_INSTALLER_TYPE} (file missing or empty)" 1>&2
            return 1
        fi
    fi

    # installing conda
    echo "Installing development environment..."
    case "${_MACH}" in
        Darwin|Linux) bash "${_INSTALLER_FILE}" -bfp "${_DEVENV}" > /dev/null ;;
                   *) cmd.exe /c "start /wait \"\" \"${_INSTALLER_FILE}\" /InstallationType=JustMe /RegisterPython=0 /AddToPath=0 /S /D=${_DEVENV} > NUL" ;;
    esac
    if ! [ $? = 0 ]; then
        echo "Error: failed to install development environment" 1>&2
        return 1
    fi
fi

# create empty env if it doesn't exist
if ! [ -d "${_ENV}" ]; then
    echo "Creating ${_NAME}..."
    if ! PYTHONPATH="" "${_BASEEXE}" create --yes --quiet "--prefix=${_ENV}" > /dev/null; then
        echo "Error: failed to create ${_NAME}" 1>&2
        return 1
    fi
fi

# check if explicitly updating or if 24 hrs since last update
if updating; then
    echo "Updating ${_NAME}..."

    if ! PYTHONPATH="" "${_BASEEXE}" update --yes --quiet --all "--prefix=${_ENV}" > /dev/null; then
        echo "Error: failed to update development environment" 1>&2
        return 1
    fi

    # set channels based on installer type
    case "${_INSTALLER_TYPE}" in
        miniconda) _CHANNEL_NAME="defaults" ;;
        miniforge) _CHANNEL_NAME="conda-forge" ;;
    esac

    if ! PYTHONPATH="" "${_BASEEXE}" install \
        --yes \
        --quiet \
        "--prefix=${_ENV}" \
        --override-channels \
        "--channel=${_CHANNEL_NAME}" \
        "--file=${_SRC}/tests/requirements.txt" \
        "--file=${_SRC}/tests/requirements-ci.txt" \
        "$([ "${_MACH}" = "Linux" ] && echo "--file=${_SRC}/tests/requirements-Linux.txt")" \
        "python=${_PYTHON}" > /dev/null; then
        echo "Error: failed to update ${_NAME}" 1>&2
        return 1
    fi

    # update timestamp
    touch "${_UPDATED}"
fi

# "install" conda
# trick conda into importing from our source code and not from site-packages
if [ -z "${PYTHONPATH+x}" ]; then
    export PYTHONPATH="${_SRC}"
else
    export PYTHONPATH="${_SRC}:${PYTHONPATH}"
fi

# copy latest shell scripts
echo "Update shell scripts..."
if ! "${_ENVEXE}" init --install > /dev/null; then
    echo "Error: failed to update shell scripts" 1>&2
    return 1
fi

# initialize conda command
echo "Initializing shell integration..."
eval "$(CONDA_AUTO_ACTIVATE=0 "${_ENVEXE}" shell.bash hook)" > /dev/null
if ! [ $? = 0 ]; then
    echo "Error: failed to initialize shell integration" 1>&2
    return 1
fi

# activate env
echo "Activating ${_NAME}..."
if ! conda activate "${_ENV}" > /dev/null; then
    echo "Error: failed to activate ${_NAME}" 1>&2
    return 1
fi

cleanup
unset -f cleanup
unset -f updating
