#!/bin/zsh

alias c='WRYN_REMOTE_CONNECTOR'
alias cmux='WRYN_REMOTE_CONNECTOR tmux'
alias sshc='WRYN_REMOTE_CONNECTOR ssh'
alias sshe='$EDITOR $WRYN_SSH_CONNECTION_LIST'
alias ssh-add-all='WRYN_SSH_ADD_CREDENTIALS'

WRYN_SSH_CONNECTION_LIST="$HOME/.config/wryn/ssh-connections"
WRYN_SSH_DELIMITER=' *| '

[ -d "$HOME/.config/wryn" ] && [ ! -f "$WRYN_SSH_CONNECTION_LIST" ] && touch "$WRYN_SSH_CONNECTION_LIST";

WRYN_SSH_ADD_CREDENTIALS() {
	local KEY_LIST=$(ssh-add -l | awk '{print $3;}')
	for key in $(grep -slR "PRIVATE" "$HOME/.ssh"); do
		[ -f "$key.pub" ] \
			&& echo $KEY_LIST | grep -q $(awk '{print $NF;}' $key.pub) \
			|| echo $KEY_LIST | grep -q $key \
			|| ssh-add "$key" >/dev/null 2>&1;
	done
}

function WRYN_REMOTE_CONNECTOR() {
	WRYN_SSH_ADD_CREDENTIALS;

	local COMMAND="$1";
	local CONNECTION="$2";

	local SSH_ID="^$CONNECTION$WRYN_SSH_DELIMITER"

	[[ $COMMAND != 'tmux' ]] && [ ! $CONNECTION ] \
		&& { echo 'Must specify a connection'; return 1 };

	grep -q "$SSH_ID" "$WRYN_SSH_CONNECTION_LIST" >/dev/null 2>&1 \
		|| [ ! $CONNECTION ] \
		|| { echo 'No such connection'; return 2 };

	case $COMMAND in;
		sshfs   ) sshfs $(WRYN_SSH_PARSE_ARGS $SSH_ID):$3 ${@:4}    ;;
		ssh     ) ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) ${@:3}         ;;
		cp-to   ) scp ${@:5} $(WRYN_SCP_PARSE_ARGS $SSH_ID $3):$4   ;;
		cp-from ) scp ${@:5} $(WRYN_SCP_PARSE_ARGS $SSH_ID):$3 $4   ;;
		copy-id ) ssh-copy-id ${@:3} $(WRYN_SSH_PARSE_ARGS $SSH_ID) ;;
		tmux    ) WRYN_REMOTE_TMUX "$SSH_ID" ${@:3}                 ;;

		* ) echo "$COMMAND not implemented" ;;
	esac
}
_WRYN_REMOTE_CONNECTOR() {
	local state;
	_arguments \
		'1: :->remote_command'\
		'2: :->connections'\
		':: :->remote_args'\
	;
	case "$state" in
		remote_command)
			compadd -X 'Choose a connection command:' \
				ssh sshfs cp-to cp-from copy-id tmux
			;;
		connections)
			compadd -X 'Choose a connection:' $(sed -n "s/$WRYN_SSH_DELIMITER.*//p" "$WRYN_SSH_CONNECTION_LIST")
			;;
		remote_args)
			_files
			case $words[2] in;
				ssh     ) _ssh;;
				sshfs   ) _sshfs;;
				cpto | cpfrom ) _scp;;
				tmux ) _tmux;;
				ssh-copy-id ) compadd $(ls);;
				* ) ;;
			esac
			;;
	esac
}
compdef _WRYN_REMOTE_CONNECTOR WRYN_REMOTE_CONNECTOR;

WRYN_SSH_PARSE_ARGS() {
	SSH_ID="$1";
	sed -n "s/$SSH_ID\(.*\)$/\1/p" "$WRYN_SSH_CONNECTION_LIST"
}

WRYN_SCP_PARSE_ARGS() {
	SSH_ID="$1";
	[ $2 ]\
		&& sed -n "s>$SSH_ID *-p \([^ ]*\)\(.*\)$>-P \1 $2\2>p" "$WRYN_SSH_CONNECTION_LIST"\
		|| sed -n "s/$SSH_ID *-p\(.*\)$/-P\1/p" "$WRYN_SSH_CONNECTION_LIST"
}

WRYN_REMOTE_TMUX() {
	local SSH_ID="$1"
	local UTF8_FLAG=$(locale charmap | grep -q 'UTF-8' && echo '-u')

	[[ "$SSH_ID" == "^$WRYN_SSH_DELIMITER" ]] \
		&& { tmux new-session -As $TMUX_DEFAULT_SESSION_NAME } \
		|| { 
			[ ! $2 ] \
				&& ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux $UTF8_FLAG new-session -As $TMUX_DEFAULT_SESSION_NAME \
				|| {
					[ $3 ] \
						&& { ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux $UTF8_FLAG ${@:2}; return 0 } \
						|| { ssh $(WRYN_SSH_PARSE_ARGS $SSH_ID) -t tmux $UTF8_FLAG new-session -As $2; };
					}
			}
}