#!/usr/bin/env bash

set -uo pipefail
shopt -s nullglob

SCRIPT_DIR="$(dirname -- "$(readlink -f -- "${BASH_SOURCE[0]}")")"
STACKS_DIR="${STACKS_DIR:-$SCRIPT_DIR}"
CONFIG_FILE="${STACKS_CONFIG:-$STACKS_DIR/manage.conf}"
LOCK_FILE="${STACKS_LOCK_FILE:-/run/lock/stacks-manage.lock}"

WHITELIST=()
BLACKLIST=()
PRIORITY=()

DISCOVERED_STACKS=()
MANAGED_STACKS=()

die() {
	echo "ERROR: $*" >&2
	exit 1
}

if [[ -f "$CONFIG_FILE" ]]; then
	bash -n "$CONFIG_FILE" || die "Syntax error in config file: $CONFIG_FILE"
	source "$CONFIG_FILE"
fi

usage() {
	cat <<EOF
Usage: $(basename "$0") COMMAND [STACK]

Commands:
  up        Create and start containers
  start     Start existing containers
  stop      Stop containers
  restart   Restart containers
  down      Stop and remove containers and networks
  pull      Pull images
  update    Pull images and apply Compose changes
  ps        Show container status
  logs      Follow logs for a stack
  config    Render Compose configuration for a stack
  images    Show configured images and versions
  list      List discovered stacks and their status
EOF
}

contains() {
	local target="$1"
	shift
	local item

	for item in "$@"; do
		[[ "$item" == "$target" ]] && return 0
	done

	return 1
}

discover_stacks() {
	local compose_file stack

	DISCOVERED_STACKS=()

	for compose_file in "$STACKS_DIR"/*/compose.yaml; do
		stack="${compose_file%/compose.yaml}"
		DISCOVERED_STACKS+=("${stack##*/}")
	done
}

is_discovered() {
	contains "$1" "${DISCOVERED_STACKS[@]}"
}

is_managed() {
	local stack="$1"

	is_discovered "$stack" || return 1
	[[ ! -e "$STACKS_DIR/$stack/.disabled" ]] || return 1
	contains "$stack" "${BLACKLIST[@]}" && return 1
	((${#WHITELIST[@]} == 0)) || contains "$stack" "${WHITELIST[@]}"
}

build_managed_stacks() {
	local stack

	MANAGED_STACKS=()

	for stack in "${PRIORITY[@]}"; do
		if is_managed "$stack" && ! contains "$stack" "${MANAGED_STACKS[@]}"; then
			MANAGED_STACKS+=("$stack")
		fi
	done

	for stack in "${DISCOVERED_STACKS[@]}"; do
		if is_managed "$stack" && ! contains "$stack" "${MANAGED_STACKS[@]}"; then
			MANAGED_STACKS+=("$stack")
		fi
	done
}

compose() {
	local stack="$1"
	shift

	(
		cd -- "$STACKS_DIR/$stack" || exit 1
		docker compose "$@"
	)
}

run_stack() {
	local stack="$1"
	local action="$2"

	echo
	echo "==> $stack: $action"

	case "$action" in
	up)
		compose "$stack" config -q &&
			compose "$stack" up -d
		;;
	update)
		compose "$stack" config -q &&
			compose "$stack" pull &&
			compose "$stack" up -d
		;;
	*)
		compose "$stack" "$action"
		;;
	esac
}

run_all() {
	local action="$1"
	local success=0 failed=0 stack i

	if [[ "$action" == "stop" || "$action" == "down" ]]; then
		for ((i = ${#MANAGED_STACKS[@]} - 1; i >= 0; i--)); do
			stack="${MANAGED_STACKS[i]}"

			if run_stack "$stack" "$action"; then
				((success += 1))
			else
				((failed += 1))
			fi
		done
	else
		for stack in "${MANAGED_STACKS[@]}"; do
			if run_stack "$stack" "$action"; then
				((success += 1))
			else
				((failed += 1))
			fi
		done
	fi

	echo
	echo "Succeeded: $success"
	echo "Failed:    $failed"

	((failed == 0))
}

stack_status() {
	local stack="$1"

	if [[ -e "$STACKS_DIR/$stack/.disabled" ]]; then
		echo "disabled"
	elif contains "$stack" "${BLACKLIST[@]}"; then
		echo "blacklisted"
	elif ((${#WHITELIST[@]} > 0)) && ! contains "$stack" "${WHITELIST[@]}"; then
		echo "not whitelisted"
	else
		echo "managed"
	fi
}

list_stacks() {
	local stack

	printf '%-24s %s\n' "STACK" "STATUS"
	printf '%-24s %s\n' "------------------------" "---------------"

	for stack in "${DISCOVERED_STACKS[@]}"; do
		printf '%-24s %s\n' "$stack" "$(stack_status "$stack")"
	done
}

stack_images() {
	local stack="$1"
	local output image

	output="$(compose "$stack" config --images)" || return 1

	while IFS= read -r image; do
		[[ -n "$image" ]] || continue
		printf '%-24s %s\n' "$stack" "$image"
	done <<<"$output"
}

list_images() {
	local target="${1:-}"
	local stack
	local failed=0

	printf '%-24s %s\n' "STACK" "IMAGE"
	printf '%-24s %s\n' "------------------------" "----------------------------------------"

	if [[ -n "$target" ]]; then
		require_target "images" "$target"
		stack_images "$target"
		return
	fi

	for stack in "${MANAGED_STACKS[@]}"; do
		stack_images "$stack" || ((failed += 1))
	done

	((failed == 0))
}

require_target() {
	local command="$1"
	local target="$2"

	[[ -n "$target" ]] || die "$command requires a stack name"
	is_discovered "$target" || die "Stack not found: $target"
	is_managed "$target" || die "Stack is excluded: $target"
}

acquire_lock() {
	mkdir -p "$(dirname "$LOCK_FILE")" ||
		die "Cannot create lock directory"
	exec {LOCK_FD}>"$LOCK_FILE" ||
		die "Cannot open lock file: $LOCK_FILE"
	flock -n "$LOCK_FD" ||
		die "Another manage.sh operation is already running"
}

discover_stacks
build_managed_stacks

ACTION="${1:-}"
TARGET="${2:-}"

[[ $# -le 2 ]] || die "Too many arguments"

case "$ACTION" in
up | start | stop | restart | down | pull | update)
	acquire_lock

	if [[ -n "$TARGET" ]]; then
		require_target "$ACTION" "$TARGET"
		run_stack "$TARGET" "$ACTION"
	else
		run_all "$ACTION"
	fi
	;;
ps)
	if [[ -n "$TARGET" ]]; then
		require_target "$ACTION" "$TARGET"
		run_stack "$TARGET" "$ACTION"
	else
		run_all "$ACTION"
	fi
	;;
logs)
	require_target "$ACTION" "$TARGET"
	compose "$TARGET" logs -f --tail=100
	;;
config)
	require_target "$ACTION" "$TARGET"
	compose "$TARGET" config
	;;
images)
	list_images "$TARGET"
	;;
list)
	[[ -z "$TARGET" ]] || die "list does not accept a stack name"
	list_stacks
	;;
help | -h | --help)
	usage
	;;
*)
	usage
	exit 1
	;;
esac