#!/bin/sh

_pyv_err() {
    printf '\033[91mError:\033[0m \033[2m%s\033[0m\n' "$1"
}

_pyv_warn() {
    printf '\033[93mWarn:\033[0m \033[2m%s\033[0m\n' "$1"
}

_pyv_log() {
    printf '\033[94mpyv:\033[0m \033[2m%s\033[0m\n' "$1"
}

_pyv_help() {
    printf 'Usage:\n  \033[92mpyv\033[0m \033[94m[ls|new|rm|go|exit]\033[0m \033[95mmyvenv\033[0m\n'
}

_pyv_avail() {
    venvs="$(for v in "$VENV_DIR"/* ; do
        if [ -d "$v" ] ; then
            if [ "$VIRTUAL_ENV" == "$v" ] ; then
                printf '\033[95m%s\033[0m, ' "*${v##*/}"
            else
                printf '\033[95m%s\033[0m, ' "${v##*/}"
            fi
        fi
        done)"
    venvs="${venvs%,*}"
    [ -n "$venvs" ] && printf 'Available venvs:\n  %s\n' "$venvs"
}

_pyv_prompt() {
    printf '%s [y/N]: ' "$1"
    read -r s
    case "$(echo "$s" | tr '[:upper:]' '[:lower:]')" in
        y|yes) return 0 ;;
        *) return 1 ;;
    esac
}

pyv() {
    VENV_OPT=
    VENV_LIST=

    PYTHON="${PYTHON:=$(command -v python 2>&-)}"
    PYTHON="${PYTHON:=$(command -v python3 2>&-)}"
    [ -z "$PYTHON" ] && _pyv_err "Could not find python in PATH" && return 1

    if ! $PYTHON -m venv -h >&- 2>&- ; then
        _pyv_err "python-venv module not available" && return 1
    fi

    if [ -z "$VENV_DIR" ] ; then
        if [ -n "$XDG_DATA_HOME" ] ; then
            VENV_DIR="$XDG_DATA_HOME/venv"
        else
            [ -z "$HOME" ] && _pyv_err "No safe venv location" && return 1
            VENV_DIR="$HOME/.local/share/venv"
        fi
    fi
    mkdir -p "$VENV_DIR"

    case "$1" in
        new|go|rm)
            VENV_OPT="$1" ; VENV_LIST="${@#${VENV_OPT}}" ; VENV_LIST="${VENV_LIST#* }" ;;
        -d)
            VENV_OPT="rm" ; VENV_LIST="${@#${VENV_OPT}}" ; VENV_LIST="${VENV_LIST#* }" ;;
        exit|quit|q)
            if command -v deactivate >&- 2>&- ; then
                _pyv_log "Deactivating venv ${VIRTUAL_ENV##*/}..."
                deactivate
            else
                _pyv_warn "No currently active venv"
            fi
            return 0 ;;
        -l*|--l*|l|ls|list)
            _pyv_avail
            return 0 ;;
        -*|help|h|"")
            _pyv_help
            avail="$(_pyv_avail)"
            [ -n "$avail" ] && printf '\n%s\n' "$avail"
            return 0 ;;
        *)
            VENV_OPT="go" ; VENV_LIST="$@" ;;
    esac

    for v in $VENV_LIST ; do
        if ! [ -d "$VENV_DIR/$v" ] && [ "$VENV_OPT" != "new" ] ; then
            _pyv_err "venv $v not found"
            VENV_LIST=
        fi
    done

    if [ -z "$VENV_LIST" ] ; then
        _pyv_help
        echo
        _pyv_avail
        return 0
    fi

    if [ "$VENV_OPT" == "new" ] ; then
        for v in $VENV_LIST ; do
            if [ -d "$VENV_DIR/$v" ] ; then
                _pyv_warn "venv $v already exists, skipping..."
            else
                _pyv_log "Creating new venv $v..."
                $PYTHON -m venv "$VENV_DIR/$v"
            fi
        done
    fi

    if [ "$VENV_OPT" == "rm" ] ; then
        for v in $VENV_LIST ; do
            if _pyv_prompt "Remove venv at '$VENV_DIR/$v'?" ; then
                _pyv_log "Removing venv $v..."
                rm -rf "$VENV_DIR/$v" && _pyv_log "Done"
            else
                _pyv_log "Keeping venv $v"
            fi
        done
    fi

    if [ "$VENV_OPT" == "go" ] ; then
        command -v deactivate >&- 2>&- && deactivate
        venv="${VENV_LIST%% *}"
        _pyv_log "Activating $venv..."

        case "${0##*/}" in
            bash|zsh|dash|ksh|sh)
                . "$VENV_DIR/$venv/bin/activate" ;;
            fish)
                source "$VENV_DIR/$venv/bin/activate.fish" ;;
            csh|tcsh)
                source "$VENV_DIR/$venv/bin/activate.csh" ;;
        esac
    fi
}