#!/bin/bash
#
# VirtualEnv shell helpers: easier create, remove, list/find and activate.
#
# Written by Radomir Stevanovic, Feb 2015 -- Feb 2017.
# Source/docs: ``https://github.com/randomir/envie``.
# Install via pip/PyPI: ``pip install envie``.

# the defaults
_ENVIE_SOURCE=$(readlink -f "${BASH_SOURCE[0]}")
_ENVIE_DEFAULT_ENVNAME=env
_ENVIE_DEFAULT_PYTHON=python
_ENVIE_CONFIG_DIR="$HOME/.config/envie"
_ENVIE_CONFIG_FILE="$_ENVIE_CONFIG_DIR/envierc"
_ENVIE_USE_DB="0"
_ENVIE_DB_PATH="$_ENVIE_CONFIG_DIR/locate.db"
_ENVIE_INDEX_ROOT="$HOME"
_ENVIE_CRON_INDEX="0"
_ENVIE_CRON_PERIOD_MIN="15"
_ENVIE_LS_INDEX="0"
_ENVIE_FIND_LIMIT="0.4"  # abort find search if takes longer, in seconds
_ENVIE_LOCATE_LIMIT="4"  # warn if index older than, in seconds
_ENVIE_UUID="28d0b2c7bc5245d5b1278015abc3f0cd"
_ENVIE_SOURCED="0"

# overwrite with values from config file
[ -r "$_ENVIE_CONFIG_FILE" ] && . "$_ENVIE_CONFIG_FILE"

# TODO: basic config options validity check (data types, ranges, etc.)

function _envie_dump_config() {
    cat <<-END
		_ENVIE_DEFAULT_ENVNAME="$_ENVIE_DEFAULT_ENVNAME"
		_ENVIE_DEFAULT_PYTHON="$_ENVIE_DEFAULT_PYTHON"
		_ENVIE_CONFIG_DIR="$_ENVIE_CONFIG_DIR"
		_ENVIE_USE_DB="$_ENVIE_USE_DB"
		_ENVIE_DB_PATH="$_ENVIE_DB_PATH"
		_ENVIE_INDEX_ROOT="$_ENVIE_INDEX_ROOT"
		_ENVIE_CRON_INDEX="$_ENVIE_CRON_INDEX"
		_ENVIE_CRON_PERIOD_MIN="$_ENVIE_CRON_PERIOD_MIN"
		_ENVIE_LS_INDEX="$_ENVIE_LS_INDEX"
		_ENVIE_FIND_LIMIT="$_ENVIE_FIND_LIMIT"
		_ENVIE_LOCATE_LIMIT="$_ENVIE_LOCATE_LIMIT"
		_ENVIE_UUID="$_ENVIE_UUID"
	END
}


#
# few generic helper functions
#

function _errmsg() {
    echo "$@" >&2
}

function _deactivate() {
    [ "$VIRTUAL_ENV" ] && deactivate
}

function _activate() {
    _deactivate
    source "$1/bin/activate"
}

function _is_virtualenv() {
    [ -e "$1/bin/activate_this.py" ] && [ -x "$1/bin/python" ]
}

# Prints all descendant of a process `ppid`, level-wise, bottom-up.
# Usage: _get_proc_descendants ppid
function _get_proc_descendants() {
    local pid ppid="$1"
    local children=$(ps hopid --ppid "$ppid")
    for pid in $children; do
        echo "$pid"
        _get_proc_descendants "$pid"
    done
}

# Kills all process trees rooted at each of the `pid`s given,
# along with all of their ancestors.
# Usage: _killtree [pid1 pid2 ...]
function _killtree() {
    while [ "$#" -gt 0 ]; do
        local pids=("$1" $(_get_proc_descendants "$1"))
        kill -TERM "${pids[@]}" &>/dev/null
        shift
    done
}

# Makes fastest temporary file: like ``mktemp``, but tries
# to create file in memory (/dev/shm) first.
function _mkftemp() {
    [ -d /dev/shm ] && mktemp --tmpdir=/dev/shm || mktemp
}

function _command_exists() {
    command -v "$1" >/dev/null 2>&1
}

function _humanized_age() {
    local age="$1"
    if ((age < 60)); then
        echo "$age second(s)"
    elif ((age < 3600)); then
        printf "~%.2g minute(s)\n" $(bc -l <<<"$age / 60")
    elif ((age < 86400)); then
        printf "~%.2g hour(s)\n" $(bc -l <<<"$age / 3600")
    else
        printf "~%.2g day(s)\n" $(bc -l <<<"$age / 86400")
    fi
}



# Creates a new environment in <path/to/env>, based on <python_exec>.
# Usage: mkenv [-p <python_exec>] [-2 | -3] [<path/to/env>] -- [options to virtualenv]
function mkenv() {
    local envpath="$_ENVIE_DEFAULT_ENVNAME" pyname="$_ENVIE_DEFAULT_PYTHON"
    local opt OPTIND pyname verbose

    while getopts ":hp:v23" opt; do
        case $opt in
            h)
                echo "Create Python (2/3) virtual environment in DEST_DIR based on PYTHON."
                echo "Usage:"
                echo "    mkenv [-2|-3|-p PYTHON] [-v] [DEST_DIR] [-- ARGS_TO_VIRTUALENV]"
                echo "    mkenv2 [-v] [DEST_DIR] [-- ARGS_TO_VIRTUALENV]"
                echo "    mkenv3 [-v] [DEST_DIR] [-- ARGS_TO_VIRTUALENV]"
                return 2;;
            p)
                pyname="$OPTARG";;
            v)
                verbose=-v;;
            2)
                pyname=python2;;
            3)
                pyname=python3;;
            \?)
                _errmsg "Invalid mkenv option: -$OPTARG. See 'mkenv -h' for help on usage."
                return 1;;
            :)
                _errmsg "Option -$OPTARG requires an argument. See 'mkenv -h' for help on usage."
                return 1;;
        esac
    done
    shift $((OPTIND-1))

    local envpath
    if [[ "$1" && "$1" != -* ]]; then
        # mkenv -2 pythonenv
        # mkenv -2 pythonenv -- xxx
        # mkenv -2 -- pythonenv -v
        envpath="$1"
        shift
        [ "$1" == -- ] && shift
    else
        # mkenv
        # mkenv -2 -v
        # mkenv -2 -- --no-pip --no-wheel
        :
    fi
    if [ -d "$envpath" ]; then
        _errmsg "Directory '$envpath' already exists."
        return 1
    fi
    echo "Creating Python virtual environment in '$envpath'."

    local pypath
    if ! pypath=$(which "$pyname"); then
        _errmsg "Python executable '$pyname' not found."
        return 1
    fi
    local pyver=$("$pypath" --version 2>&1)
    if [[ ! $pyver =~ Python ]]; then
        _errmsg "Unrecognized Python version of executable: '$pypath'."
        return 1
    fi
    echo "Using $pyver ($pypath)."

    mkdir -p "$envpath"
    cd "$envpath"
    local output
    output=$(virtualenv --no-site-packages -p "$pypath" $verbose "$@" . 2>&1)
    if [ $? -ne 0 ]; then
        _errmsg "$output"
    else
        [ "$verbose" ] && echo "$output"
        _activate .
    fi
    cd - >/dev/null
}

function mkenv2() {
    mkenv -2 "$@"
}

function mkenv3() {
    mkenv -3 "$@"
}


# Destroys the active environment.
# Usage (while env active): rmenv
function rmenv() {
    local envpath="$VIRTUAL_ENV"
    if [ ! "$envpath" ]; then
        _errmsg "Active virtual environment not detected."
        return 1
    fi
    deactivate
    if _is_virtualenv "$envpath"; then
        rm -rf "$envpath"
    else
        _errmsg "Invalid VirtualEnv path in VIRTUAL_ENV: '$envpath'."
        return 1
    fi
}


# Lists all environments below the <start_dir>.
# Usage: lsenv [<start_dir> [<avoid_subdir>]]
function _lsenv_find() {
    local dir="${1:-.}" avoid="${2:-}"
    find "$dir" -path "$avoid" -prune -o \
        -name .git -o -name .hg -o -name .svn -prune -o -path '*/bin/activate_this.py' \
        -exec dirname '{}' \; 2>/dev/null | xargs -d'\n' -n1 -r dirname
}

# `lsenv` via `locate`
# Compatible with: lsenv [<start_dir> [<avoid_subdir>]]
function _lsenv_locate() {
    local dir="${1:-.}" avoid="${2:-}"
    local absdir=$(readlink -e "$dir")
    [ "$absdir" = / ] && dir=/
    if (( _ENVIE_LS_INDEX && $(_envie_db_age) > _ENVIE_LOCATE_LIMIT )); then
        __envie_updatedb
    fi
    locate -d "$_ENVIE_DB_PATH" --existing "$absdir"'*/bin/activate_this.py' \
        | sed -e 's#/bin/activate_this\.py$##' -e "s#^$absdir#$dir#"
}

# Run `lsenv` via both `find` and `locate` in parallel and:
# - wait `$_ENVIE_FIND_LIMIT` seconds for `find` to finish
# - if it finishes on time, take those results, as they are the most current and accurate
# - if find takes longer, kill it and wait for `locate` results
function _lsenv_locate_vs_find_race() {
    set +m
    local p_pid_find=$(_mkftemp) p_pid_locate=$(_mkftemp) p_pid_timer=$(_mkftemp)
    local p_ret_find=$(_mkftemp) p_ret_locate=$(_mkftemp)
    { __find_and_return "$@" & echo $! >"$p_pid_find"; } 2>/dev/null
    { __locate_and_return "$@" & echo $! >"$p_pid_locate"; } 2>/dev/null
    { __find_fast_bailout & echo $! >"$p_pid_timer"; } 2>/dev/null
    wait
    if [ -e "$p_ret_find" ]; then
        cat "$p_ret_find"
    elif [ -e "$p_ret_locate" ]; then
        cat "$p_ret_locate"
        local db_age=$(_envie_db_age)
        if (( db_age > _ENVIE_LOCATE_LIMIT )); then
            _errmsg "NOTE: results are based on a db from $(_humanized_age "$db_age") ago, and may not include all current virtualenvs."
            _errmsg "Use 'lsenv -f' to force manual search, or run 'envie update' to update the database."
        fi
    fi
    rm -f "$p_pid_find" "$p_pid_locate" "$p_pid_timer" "$p_ret_find" "$p_ret_locate"
    set -m
}
function __find_and_return() {
    _lsenv_find "$@" >"$p_ret_find"
    _killtree $(<"$p_pid_locate") $(<"$p_pid_timer")
    rm -f "$p_ret_locate"
}
function __locate_and_return() {
    _lsenv_locate "$@" >"$p_ret_locate"
}
function __find_fast_bailout() {
    sleep "$_ENVIE_FIND_LIMIT"
    _killtree $(<"$p_pid_find")
    rm -f "$p_ret_find"
}


# Lists all virtual environments below `start_dir` using `find` or `locate`, if 
# specified explicitly (via opts or env var). If search method undefined, use
# `find vs locate race` where find is allowed to run for 400ms, and then aborted
# favor of locate.
# Usage: lsenv [-f|-l] [<start_dir> [<avoid_subdir>]]
function lsenv() {
    if [[ "$1" == "-f" ]]; then
        shift
        _lsenv_find "$@"
    elif [[ "$1" == "-l" ]]; then
        shift
        _lsenv_locate "$@"
    elif (( ! _ENVIE_USE_DB )); then
        _lsenv_find "$@"
    else
        _lsenv_locate_vs_find_race "$@"
    fi
}


# Finds the closest env by first looking down and then dir-by-dir up the tree.
function lsupenv() {
    local list len=0 dir=. prevdir quiet=0

    if [[ "$1" == "-q" ]]; then
        quiet=1
        shift
    fi

    while [ "$len" -eq 0 ] && [ "$(readlink -e "$prevdir")" != / ]; do
        if ((quiet)); then
            list=$(lsenv "$dir" "$prevdir" 2>/dev/null)
        else
            list=$(lsenv "$dir" "$prevdir")
        fi
        [ "$list" ] && len=$(wc -l <<<"$list") || len=0
        prevdir="$dir"
        dir="$dir/.."
    done
    echo "$list"
}


# Changes virtual environment (activates the closest env if unique, or asks to
# select among the closest envs)
function chenv() {
    local IFS envlist env len=0 quiet=0

    if [[ "$1" == "-q" ]]; then
        quiet=1
        shift
    fi

    envlist=$(lsupenv -q)
    [ "$envlist" ] && len=$(wc -l <<<"$envlist")
    if [ "$len" -eq 1 ]; then
        _activate "$envlist"
    elif [ "$len" -eq 0 ]; then
        (( ! quiet )) && _errmsg "No environments found."
        return 1
    else
        (( quiet )) && return 1
        IFS=$'\n'
        select env in $envlist; do
            if [ "$env" ]; then
                _activate "$env"
                break
            fi
        done
    fi
}


# `cd` to active env base dir, or fail if no env active
function cdenv() {
    if [ "$VIRTUAL_ENV" ]; then
        cd "$VIRTUAL_ENV"
    else
        _errmsg "Virtual environment not active. Use 'chenv' to activate."
        return 1
    fi
}


# faster envie, using locate

function _envie_db_age() {
    [ ! -e "$_ENVIE_DB_PATH" ] && return 1
    echo $(( $(date +%s) - $(date -r "$_ENVIE_DB_PATH" +%s) ))
}

function _envie_locate_exists() {
    if ! _command_exists locate || ! _command_exists updatedb; then
        _errmsg "locate/updatedb not installed. Failing-back to find."
        return 1
    fi
}


function __envie_initdb() {
    _envie_locate_exists || return 1
    echo -n "Indexing environments in '$_ENVIE_INDEX_ROOT'..."
    __envie_updatedb
    echo "Done."
}

function __envie_updatedb() {
    updatedb -l 0 -o "$_ENVIE_DB_PATH" -U "$_ENVIE_INDEX_ROOT" \
        --prune-bind-mounts 0 \
        --prunenames ".git .bzr .hg .svn" \
        --prunepaths "/var/spool /media /home/.ecryptfs /var/lib/schroot" \
        --prunefs "NFS nfs nfs4 rpc_pipefs afs binfmt_misc proc smbfs autofs \
                   iso9660 ncpfs coda devpts ftpfs devfs mfs shfs sysfs cifs \
                   lustre tmpfs usbfs udf fuse.glusterfs fuse.sshfs curlftpfs \
                   ecryptfs fusesmb devtmpfs"
}

# Add to .bashrc
function __envie_register() {
    mkdir -p "$_ENVIE_CONFIG_DIR"
    
    local bashrc=~/.bashrc
    [ ! -w "$bashrc" ] && _errmsg "$bashrc not writeable." && return 1

    [ -z "$_ENVIE_SOURCE" ] && _errmsg "Envie source script not found." && return 2

    if grep "$_ENVIE_UUID" "$bashrc" &>/dev/null; then
        _errmsg "Envie already registered in $bashrc."
        return
    fi

    cat >>"$bashrc" <<-END
		# Load 'envie' (Python VirtualEnv helpers)  #$_ENVIE_UUID
		[ -f "$_ENVIE_SOURCE" ] && source "$_ENVIE_SOURCE"  #$_ENVIE_UUID
	END
    echo "Envie added to $bashrc."
}

# Remove from .bashrc
function __envie_unregister() {
    local bashrc=~/.bashrc
    [ ! -w "$bashrc" ] && _errmsg "$bashrc not writeable." && return 1

    if ! cp -a "$bashrc" "$_ENVIE_CONFIG_DIR/.bashrc.backup"; then
        _errmsg "Failed to backup $bashrc before modifying."
        return 1
    fi
    if sed -e "/$_ENVIE_UUID/d" "$bashrc" -i; then
        echo "Envie removed from $bashrc."
    fi
}

function __envie_configure() {
    local ans

    read -p "Use locate/updatedb for faster search [Y/n]? " ans
    case "$ans" in
        N|n) _ENVIE_USE_DB=0;;
        *)
            if _envie_locate_exists; then
                _ENVIE_USE_DB=1
            else
                _ENVIE_USE_DB=0
            fi;;
    esac

    if (( _ENVIE_USE_DB )); then
        read -p "Common ancestor dir of all environments to be indexed [$_ENVIE_INDEX_ROOT]: " ans
        if [ "$ans" ]; then
            if [ -d "$ans" ]; then
                _ENVIE_INDEX_ROOT=$(readlink -f "$ans")
            else
                echo "Invalid dir $ans. Skipping."
            fi
        fi

        read -p "Update index periodically (every ${_ENVIE_CRON_PERIOD_MIN}min) [Y/n]? " ans
        case "$ans" in
            N|n) _ENVIE_CRON_INDEX=0;;
            *) _ENVIE_CRON_INDEX=1;;
        esac

        read -p "Refresh stale index before each search [Y/n]? " ans
        case "$ans" in
            N|n) _ENVIE_LS_INDEX=0;;
            *) _ENVIE_LS_INDEX=1;;
        esac
    else
        _ENVIE_CRON_INDEX=0
        _ENVIE_LS_INDEX=0
    fi

    read -p "Add to ~/.bashrc [Y/n]? " ans
    case "$ans" in
        N|n) ;;
        *) __envie_register;;
    esac

    # save config
    mkdir -p "$_ENVIE_CONFIG_DIR"
    _envie_dump_config >"$_ENVIE_CONFIG_FILE" && echo "Config file written to $_ENVIE_CONFIG_FILE."

    # add/remove from crontab
    _envie_update_crontab && echo "Crontab updated."

    # db (re-)index
    (( _ENVIE_USE_DB )) && __envie_initdb
}

function _envie_update_crontab() {
    if (( _ENVIE_USE_DB && _ENVIE_CRON_INDEX )); then
        # add
        (
            crontab -l | grep -v "$_ENVIE_UUID"
            echo "*/${_ENVIE_CRON_PERIOD_MIN} * * * * $_ENVIE_SOURCE update  #$_ENVIE_UUID"
        ) 2>/dev/null | crontab -
    else
        # remove
        crontab -l | grep -v "$_ENVIE_UUID" | crontab -
    fi
}

function _envie_exec() {
    if (( _ENVIE_SOURCED )); then
        eval "$@"
    else
        exec "$@"
    fi
}

# main -- handle direct call: ``envie <cmd> <args> | <script>``
# AND act as an chenv alias
function __envie_main() {
    local cmd script;

    if [ $# -gt 0 ]; then
        if [ -f "$1" ]; then
            cmd=python
        else
            cmd="$1"
            shift
        fi
    fi

    case "$cmd" in
        reg|register)
            __envie_register;;
        unreg|unregister)
            __envie_unregister;;
        init|initdb)
            __envie_initdb;;
        update|updatedb)
            __envie_updatedb;;
        config|configure)
            __envie_configure;;
        exec|run)
            # execute CMD
            script="$1"
            shift
            if [[ $(type -t "$script") =~ alias|function|builtin|file ]]; then
                (
                    # move closer to script for env detection
                    if [ -f "$script" ]; then
                        script=$(readlink -f "$script")
                        cd "$(dirname "$script")"
                    fi
                    chenv -q
                    _envie_exec "$script" "$@"
                )
            else
                _errmsg "'$script' is not a valid command. See 'envie --help'."
                return 1
            fi
            ;;
        python)
            # run Python SCRIPT in current shell,
            # or just run Python in interactive mode if SCRIPT missing
            script="$1"
            shift
            # handle: `envie python`
            if [ ! "$script" ]; then
                chenv -q
                _envie_exec python
                return "$?"
            fi
            # handle: `envie python non-existing-file`
            if [ ! -f "$script" ]; then
                _errmsg "'$script' is not a file. See 'envie --help'."
                return 1
            fi
            # handle: `envie python path/to/some/script.py`
            (
                # move closer to script for proper env detection
                script=$(readlink -f "$script")
                cd "$(dirname "$script")"
                chenv -q
                _envie_exec python "$script" "$@"
            )
            ;;
        create)
            mkenv "$@";;
        remove)
            rmenv "$@";;
        list)
            lsenv "$@";;
        go)
            chenv "$@";;
        help|--help|*)
            __envie_usage;;
    esac
}

function __envie_usage() {
    cat <<-END
		Usage:
		    envie {python SCRIPT | exec CMD | config | register | unregister | init | update}
		    envie SCRIPT
		
		Commands:
		    python SCRIPT run Python SCRIPT in the closest environment
		    exec CMD      execute CMD in the closest environment. CMD can be a
		                  script file, command, builtin, alias, or a function.
		
		    config        interactively configure envie
		    init          index virtualenvs below $_ENVIE_INDEX_ROOT
		    update        update index
		    register      add envie to .bashrc
		    unregister    remove envie from .bashrc
		    help          this help
		
		    create [ENV]  alias for mkenv
		    remove        alias for rmenv
		    list [DIR]    alias for lsenv
		    go            alias for chenv
		
		The second form is a shorthand for executing python scripts in the closest 
		virtual environment, without the need for manual env activation. It's convenient
		for hash bangs:
		    #!/usr/bin/env envie
		    # Python script here will be executed in the closest virtual env
		
		Examples:
		    envie python              # run interactive Python shell in the closest env
		    envie manage.py shell     # run Django shell in the project env (auto activate)
		    envie exec /path/to/executable    # execute an executable in the closest env
	END
}


# bash/readline completions
function __envie_complete() {
    local word="${COMP_WORDS[COMP_CWORD]}"
    local prev="${COMP_WORDS[COMP_CWORD-1]}"
    local cmds="python exec config register unregister init update help create remove list go"

    case "$prev" in
        exec)
            COMPREPLY=($(compgen -A alias -A builtin -A command -A file -A function -- $word))
            ;;
        python)
            COMPREPLY=($(compgen -A file -- $word))
            ;;
        *)
            COMPREPLY=($(compgen -W "$cmds" -A file -- $word))
            ;;
    esac
}

complete -F __envie_complete envie

# mask executable envie with function to be able to for eg. activate envs
function envie() {
    _ENVIE_SOURCED="1"
    __envie_main "$@"
}

[ $# -gt 0 ] && __envie_main "$@"
