#!/bin/zsh

source ${$(realpath "$0"):h:h}/lib/common.zsh

zparseopts -D -E \
	c=cdvenv -cdvenv=cdvenv \
	h=help -help=help \
	n=newvenv -new=newvenv \
	t=tmsu -tmsu=tmsu \
	p=pyenv -pyenv=pyenv \
	f=find -find=find

if [[ $help ]]; then
	<<-HELP
	Usage: ${0:t} [-c] [DIR]
	       ${0:t} -t

	If DIR contains bin/activate, start a shell with that virtualenv.
	Otherwise recursively look up virtualenvs in DIR (by looking for dirs ending
	in "venv") and select one with fzf.

	OPTIONS:
	 -c, --cdvenv  Change to the venv directory instead of going back to the 
				   current working directory (not for pyenv virtualenvs)
	 -t, --tmsu    list only venvs named "venv" in a directory with tmsu tag lang=python
	 -p, --pyenv   Use pyenv virtualenvs
	 -f, --find    Search for virtualenvs by looking for directories ending with 'venv'

	 If neither -p or -f are given, the default behaviour is to use both
	HELP
	exit
fi

ORIGDIR=$PWD

venvsh() {
	cd $1
	. bin/activate
	if [[ -z $cdvenv ]]; then
		cd $ORIGDIR
	fi
	$SHELL
}

pyenvsh() {
	export PYENV_VIRTUALENV_DISABLE_PROMPT=1
	eval "$(pyenv init -)"
	eval "$(pyenv virtualenv-init -)"
	pyenv activate "$1"
	$SHELL
}

if [[ -n $tmsu ]]; then
	venv=$(
		for i in $(tmsu files --directory lang=python); do
			[[ -d $i/venv ]] && echo $i/venv
		done | fzf)
	if [[ -n $venv ]]; then
		exec venvsh $venv
	else
		exit 1
	fi
fi

if [[ $newvenv ]]; then
	python -mvenv ${1:-venv}
	exec venvsh ${1:-venv}
fi

if [[ -n $1 && -e $1/bin/activate ]]; then
	exec venvsh $1
fi

pyenvs() {
	if exists pyenv && pyenv commands | grep -q 'virtualenvs'; then
		pyenv virtualenvs --bare --skip-aliases | sed -e 's/^/pyenv!/'
	fi
}

if [[ $pyenv && ! $find ]]; then
	venv=$(pyenvs | fzf)
elif [[ ! $pyenv && $find ]]; then
	venv=$(fd -I -t d 'venv$' $1 | fzf)
else
	venv=$(cat <(fd -I -t d 'venv$' $1) <(pyenvs) | fzf)
fi
if [[ -n $venv ]]; then
	if [[ $venv =~ pyenv\!* ]]; then
		exec pyenvsh ${venv##pyenv!}
	else
		exec venvsh $venv
	fi
fi