#!/bin/bash
#
# VirtualEnv shell helpers: easier create, remove, list/find and activate.
# Written by Radomir Stevanovic, Feb 2015.

_ENVIE_SOURCE=$(readlink -f "${BASH_SOURCE[0]}")
_ENVIE_DEFAULT_ENVNAME=env
_ENVIE_DEFAULT_PYTHON=python
_ENVIE_CONFIG_DIR="$HOME/.config/envie"
_ENVIE_DB_PATH="$_ENVIE_CONFIG_DIR/locate.db"
_ENVIE_INDEX_ROOT="$HOME"
_ENVIE_FIND_LIMIT=0.1  # in seconds
_ENVIE_UUID="28d0b2c7bc5245d5b1278015abc3f0cd"

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

# Creates a new environment in <path/to/env>, based on <python_exec>.
# Usage: mkenv [<path/to/env>] [<python_exec>]
function mkenv() {
    local envpath="${1:-$_ENVIE_DEFAULT_ENVNAME}" output
    if [ -d "$envpath" ]; then
        fail "Directory '$envpath' already exists."
        return 1
    fi
    echo "Creating python environment in '$envpath'."

    local pyexe="${2:-$_ENVIE_DEFAULT_PYTHON}" pypath
    if ! pypath=$(which "$pyexe"); then
        fail "Python executable '$pyexe' not found, failing-back to: '$_ENVIE_DEFAULT_PYTHON'."
        pypath=$(which "$_ENVIE_DEFAULT_PYTHON")
    fi
    local pyver=$("$pypath" --version 2>&1)
    if [[ ! $pyver =~ Python ]]; then
        fail "Unrecognized Python version/executable: '$pypath'."
        return 1
    fi
    echo "Using $pyver ($pypath)."

    mkdir -p "$envpath"
    cd "$envpath"
    output=$(virtualenv --no-site-packages -p "$pypath" -v . 2>&1)
    if [ $? -ne 0 ]; then
        echo -en "$output"
    fi
    . bin/activate
    cd - >/dev/null
}

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

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" ]
}

# 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")
    locate -d "$_ENVIE_DB_PATH" "$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"
    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"
    __kill $(<"$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"
    __kill "$(<"$p_pid_find")"
    rm -f "$p_ret_find"
}
function __kill() {
    while [ "$#" -gt 0 ]; do
        _kill_proc_tree "$1" &>/dev/null
        shift
    done
}

# 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 a complete process tree rooted at `pid`.
# Usage: _kill_proc_tree pid
function _kill_proc_tree() {
    local pids pids_list="$(_get_proc_descendants "$1") $1"
    IFS=' ' pids=($pids_list)
    command kill -TERM "${pids[@]}"
}

# Make 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 lsenv() {
    _db_exists && _lsenv_locate_vs_find_race "$@" || _lsenv_find "$@"
}

# Finds the closest env by first looking down and then dir-by-dir up the tree.
function lsupenv() {
    local list len=0 dir=. prevdir
    while [ "$len" -eq 0 ] && [ "$(readlink -e "$prevdir")" != / ]; do
        list=$(lsenv "$dir" "$prevdir")
        [ "$list" ] && len=$(wc -l <<<"$list") || len=0
        prevdir="$dir"
        dir="$dir/.."
    done
    echo "$list"
}

function cdenv() {
    local OLDIFS envlist env len=0

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

alias wkenv=cdenv


# faster envie, using locate

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

function _db_exists() {
    [ -e "$_ENVIE_DB_PATH" ]
}

function __envie_initdb() {
    if ! _command_exists locate || ! _command_exists updatedb; then
        fail "locate/updatedb not installed. Failing-back to find."
        return 1
    fi
    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"
}

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

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

    if grep "$_ENVIE_UUID" "$bashrc" &>/dev/null; then
        fail "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" ] && fail "$bashrc not writeable." && return 1

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


# main -- handle direct call: ``envie <cmd> <args>``
function __envie_main() {
    local cmd;

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

    case "$cmd" in
        reg|register)
            __envie_register;;
        unreg|unregister)
            __envie_unregister;;
        init|initdb)
            __envie_initdb;;
        update|updatedb)
            __envie_updatedb;;
        *)
            __envie_usage;;
    esac
}

function __envie_usage() {
	cat <<-END
		Usage:
		    envie {register | unregister | init | update}

		Commands:
		    register      add envie to .bashrc
		    unregister    remove envie from .bashrc
		    init          index virtualenvs below $HOME
		    update        update index
	END
}

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