#!/bin/zsh

alias c='WRYN_REMOTE_CONNECTOR'
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() {
	for key in $(grep -slR "PRIVATE" "$HOME/.ssh"); do
		ssh-add "$key" && { echo "Added $key" } || { echo "Failed to add $key" };
	done
}

function WRYN_REMOTE_CONNECTOR() {
	local COMMAND="$1";
	local CONNECTION="$2";

	local SSH_ID="$CONNECTION$WRYN_SSH_DELIMITER"

	[ $2 ] \
		|| { echo 'Must specify a connection'; return 1 };
	grep -q "$SSH_ID" "$WRYN_SSH_CONNECTION_LIST" >/dev/null 2>&1 \
		|| { 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 $4):$3   ;;
		cp-from ) scp ${@:5} $(WRYN_SCP_PARSE_ARGS $SSH_ID):$3 $4   ;;
		copy-id ) ssh-copy-id $(WRYN_SSH_PARSE_ARGS $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
			;;
		connections)
			compadd -X 'Choose a conection:' $(sed "s/$WRYN_SSH_DELIMITER.*//" "$WRYN_SSH_CONNECTION_LIST")
			;;
		remote_args)
			_files
			case $words[2] in;
				ssh     ) _ssh;;
				sshfs   ) _sshfs;;
				cpto | cpfrom ) _scp;;
				copy-id ) _ssh-copy-id;;
				* ) ;;
			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"
}