#!/usr/bin/env bash
set -euo pipefail

CONFIG_DIR="$HOME/.config/mtm-ssh-menu"
USE_KNOWN_HOSTS=0
USE_KNOWN_HOSTS_USER="root"

usage() {
	cat <<'EOF'
Usage: sshm [--config-dir DIR] [-k [user]]

Options:
  -h | --help                Show this help message
  --config-dir DIR           Use a custom config directory
	-k | --known-hosts [user]  Include known_hosts

The config directory must contain:
  - global.yaml
  - hosts/*.yaml
EOF
}

parse_args() {
	while [[ $# -gt 0 ]]; do
		case "$1" in
			--config-dir)
				if [[ $# -lt 2 ]]; then
					printf 'Error: --config-dir requires a value.\n' >&2
					exit 1
				fi
				CONFIG_DIR="$2"
				shift 2
				;;
			--known-hosts|-k)
				USE_KNOWN_HOSTS=1
				if [ $# -eq 1 ]; then
					USE_KNOWN_HOSTS_USER="root"
					shift
				else
					USE_KNOWN_HOSTS_USER="$2"
					shift 2
				fi
				;;
			--help|-h)
				usage
				exit 0
				;;
			*)
				printf 'Error: unknown argument: %s\n' "$1" >&2
				usage >&2
				exit 1
				;;
		esac
	done
}
parse_args "$@"

dependency_check() {
	if ! command -v yq >/dev/null 2>&1; then
		printf 'Error: yq is required but not installed.\n' >&2
		exit 1
	fi
	if ! command -v jq >/dev/null 2>&1; then
		printf 'Error: jq is required but not installed.\n' >&2
		exit 1
	fi
	if ! command -v ssh >/dev/null 2>&1; then
		printf 'Error: ssh is required but not installed.\n' >&2
		exit 1
	fi
	if ! command -v fzf >/dev/null 2>&1; then
		printf 'Error: fzf is required but not installed.\n' >&2
		exit 1
	fi
}
dependency_check

GLOBAL_CONFIG="$CONFIG_DIR/global.yaml"
HOSTS_DIR="$CONFIG_DIR/hosts"
DEFAULT_SSH_USER="root"
DEFAULT_SSH_PORT="22"
DEFAULT_SSH_OPTIONS=""
SSH_JUMP_HOSTS={}
SERVERS=""

load_config() {
	local GLOBAL_CONFIG_CONTENT
	GLOBAL_CONFIG_CONTENT="$(<"$GLOBAL_CONFIG")"
	DEFAULT_SSH_USER=$(yq -r '.ssh.default_user // "'$DEFAULT_SSH_USER'"' <<<"$GLOBAL_CONFIG_CONTENT")
	DEFAULT_SSH_PORT=$(yq -r '.ssh.default_port // "'$DEFAULT_SSH_PORT'"' <<<"$GLOBAL_CONFIG_CONTENT")
	DEFAULT_SSH_OPTIONS=$(yq -r '.ssh.default_options // "'"$DEFAULT_SSH_OPTIONS"'"' <<<"$GLOBAL_CONFIG_CONTENT")
	SSH_JUMP_HOSTS=$(yq -r '.ssh.jump_hosts // {}' <<<"$GLOBAL_CONFIG_CONTENT")

	shopt -s nullglob
	for file in "$HOSTS_DIR"/*.yaml; do
		local FILE_CONTENT GROUP_NAME GROUP_SERVERS INDEX
		FILE_CONTENT="$(<"$file")"
		GROUP_NAME=$(basename "$file")
		GROUP_NAME=${GROUP_NAME%.*}
		GROUP_NAME=$(yq -r '.group // "'"$GROUP_NAME"'"' <<<"$FILE_CONTENT")
		GROUP_SERVERS=$(yq -r '.servers // []' <<<"$FILE_CONTENT")
		INDEX=0
		
		for ((i=0; i<$(jq -r '. | length' <<<"$GROUP_SERVERS"); i++)); do
			local SSH_SERVER_NAME SSH_SERVER_ALIASES ALIASES SSH_SERVER_HOST SSH_SERVER_PORT SSH_SERVER_USER
			SSH_SERVER_NAME=$(jq -r '.'"[$i]"'.name // ""' <<<"$GROUP_SERVERS")
			SSH_SERVER_ALIASES=$(jq -r '.'"[$i]"'.aliases // ""' <<<"$GROUP_SERVERS")
			SSH_SERVER_HOST=$(jq -r '.'"[$i]"'.host // ""' <<<"$GROUP_SERVERS")
			SSH_SERVER_PORT=$(jq -r '.'"[$i]"'.port // "'"$DEFAULT_SSH_PORT"'"' <<<"$GROUP_SERVERS")
			SSH_SERVER_USER=$(jq -r '.'"[$i]"'.user // "'"$DEFAULT_SSH_USER"'"' <<<"$GROUP_SERVERS")
			if [ "$SSH_SERVER_ALIASES" != "" ]; then
				# shellcheck disable=SC2027
				ALIASES="("$(echo "$SSH_SERVER_ALIASES" | jq -r 'join(", ")')")"
			fi
			if [ "$SERVERS" != "" ]; then
				SERVERS+="\n"
			fi
			SERVERS+="${GROUP_NAME} | ${INDEX} | ${SSH_SERVER_HOST} | ${SSH_SERVER_NAME} ${ALIASES} | ${SSH_SERVER_USER} | ${SSH_SERVER_PORT}"
			INDEX=$((INDEX+1))
		done

	done
	shopt -u nullglob
}
load_known_hosts() {
	local FILE_CONTENT GROUP_NAME
	FILE_CONTENT=$(cat "$HOME"/.ssh/known_hosts | awk -F ' ' '{print $1}' | sort | awk '!seen[$1]++')
	GROUP_NAME="Local"
	for host in $FILE_CONTENT; do
		if [ "$SERVERS" != "" ]; then
			SERVERS+="\n"
		fi
		SERVERS+="${GROUP_NAME} | ${host} | ${host} | ${host} | ${USE_KNOWN_HOSTS_USER} | 22"
	done
}
popup_menu() {
	local SERVER
	SERVER=$(echo -e "${SERVERS}" | column -t -s "|" -o "|" | fzf -e --layout=reverse --with-nth=1,3,4,5,6 --delimiter="|")
	local GROUP_NAME HOST_INDEX
	GROUP_NAME=$(echo "$SERVER" | awk -F '|' '{print $1}' | xargs)
	HOST_INDEX=$(echo "$SERVER" | awk -F '|' '{print $2}' | xargs)
	ssh_connect "$GROUP_NAME" "$HOST_INDEX"
}
ssh_connect() {
	if [[ "$HOST_INDEX" != *"."* ]]; then
		local SERVER SSH_SERVER_USER SSH_SERVER_HOST SSH_SERVER_PORT SSH_SERVER_OPTIONS SSH_JUMP_HOST
		SERVER=$(cat "$HOSTS_DIR/$1.yaml" | yq -r '.servers'"[$2]")
		SSH_SERVER_USER=$(jq -r '.user // "'"$DEFAULT_SSH_USER"'"' <<<"$SERVER")
		SSH_SERVER_HOST=$(jq -r '.host // ""' <<<"$SERVER")
		SSH_SERVER_PORT=$(jq -r '.port // "'"$DEFAULT_SSH_PORT"'"' <<<"$SERVER")

		SSH_SERVER_OPTIONS="-p $SSH_SERVER_PORT"
		SSH_JUMP_HOST=$(jq -r '.jump_host // ""' <<<"$SERVER")
		if [ "$SSH_JUMP_HOST" != "" ]; then
			SSH_JUMP_HOST=$(jq -r '.'"$SSH_JUMP_HOST"' // ""' <<<"${SSH_JUMP_HOSTS}")
			SSH_JUMP_HOST="-t $SSH_JUMP_HOST ssh -p $SSH_SERVER_PORT"
			SSH_SERVER_OPTIONS=""
		fi

		# shellcheck disable=SC2086
		ssh ${SSH_SERVER_OPTIONS} ${SSH_JUMP_HOST} "${SSH_SERVER_USER}"@"${SSH_SERVER_HOST}"
	else
		ssh "${USE_KNOWN_HOSTS_USER}"@"${2}"
	fi
}

main() {
	load_config
	if [ "$USE_KNOWN_HOSTS" == 1 ]; then
		load_known_hosts
	fi
	popup_menu
}

main
