#!/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}"

DEFAULT_MAX_PARALLEL=2

CONFIG_GROUP_NAMES=()
CONFIG_GROUP_LIMITS=()
CONFIG_GROUP_STACKS=()
CONFIGURED_STACKS=()
EXCLUDED_STACKS=()

DISCOVERED_STACKS=()
MANAGED_STACKS=()
GROUP_NAMES=()
GROUP_LIMITS=()
GROUP_STACKS=()
UNASSIGNED_STACKS=()

RUN_SUCCESS=0
RUN_FAILED=0
RUN_FAILED_STACKS=()
RUN_SKIPPED_GROUPS=()

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

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 group/status

Bulk execution:
  up/start/restart/pull/update run groups in declaration order.
  stop/down run groups in reverse order.
  Stacks within a group run in parallel up to that group's limit.
  ps/images/list keep ordered output.

Configuration: $CONFIG_FILE
EOF
}

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

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

	return 1
}

validate_parallel_limit() {
	local value="$1"

	[[ "$value" =~ ^[0-9]+$ ]] ||
		die "parallel limit must be a non-negative integer: $value"
}

validate_name() {
	local kind="$1"
	local value="$2"

	[[ "$value" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]*$ ]] ||
		die "invalid $kind name: $value"
}

# Set the default concurrency for groups without an override.
max_parallel() {
	[[ $# -eq 1 ]] || die "max_parallel expects exactly one value"
	validate_parallel_limit "$1"
	DEFAULT_MAX_PARALLEL="$1"
}

# Declare an ordered group and its stacks.
# Usage: group NAME [--parallel N] STACK...
group() {
	[[ $# -ge 2 ]] || die "group expects NAME [--parallel N] STACK..."

	local name="$1"
	local limit=""
	local stack
	shift

	validate_name "group" "$name"
	[[ "$name" != "unassigned" ]] ||
		die "group name is reserved: $name"
	contains "$name" "${CONFIG_GROUP_NAMES[@]}" &&
		die "group is declared more than once: $name"

	if [[ "${1:-}" == "--parallel" ]]; then
		[[ $# -ge 3 ]] || die "group $name: --parallel expects a value and at least one stack"
		limit="$2"
		validate_parallel_limit "$limit"
		shift 2
	fi

	[[ $# -ge 1 ]] || die "group has no stacks: $name"

	for stack in "$@"; do
		validate_name "stack" "$stack"
		contains "$stack" "${CONFIGURED_STACKS[@]}" &&
			die "stack is assigned to more than one group: $stack"
		contains "$stack" "${EXCLUDED_STACKS[@]}" &&
			die "stack is both grouped and excluded: $stack"
		CONFIGURED_STACKS+=("$stack")
	done

	CONFIG_GROUP_NAMES+=("$name")
	CONFIG_GROUP_LIMITS+=("$limit")
	CONFIG_GROUP_STACKS+=("$*")
}

# Exclude stacks from all bulk operations.
exclude() {
	[[ $# -ge 1 ]] || die "exclude expects at least one stack"

	local stack
	for stack in "$@"; do
		validate_name "stack" "$stack"
		contains "$stack" "${EXCLUDED_STACKS[@]}" &&
			die "stack is excluded more than once: $stack"
		contains "$stack" "${CONFIGURED_STACKS[@]}" &&
			die "stack is both grouped and excluded: $stack"
		EXCLUDED_STACKS+=("$stack")
	done
}

load_config() {
	[[ -f "$CONFIG_FILE" ]] || return 0

	# The trusted config calls max_parallel, group, and exclude above.
	source "$CONFIG_FILE" || die "cannot load config file: $CONFIG_FILE"
}

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() {
	is_discovered "$1" && ! contains "$1" "${EXCLUDED_STACKS[@]}"
}

validate_config_stacks() {
	local stack

	for stack in "${CONFIGURED_STACKS[@]}"; do
		is_discovered "$stack" || die "configured stack not found: $stack"
	done

	for stack in "${EXCLUDED_STACKS[@]}"; do
		is_discovered "$stack" || die "excluded stack not found: $stack"
	done
}

build_execution_groups() {
	local i stack limit
	local -a stacks=()

	GROUP_NAMES=()
	GROUP_LIMITS=()
	GROUP_STACKS=()
	MANAGED_STACKS=()
	UNASSIGNED_STACKS=()

	for ((i = 0; i < ${#CONFIG_GROUP_NAMES[@]}; i++)); do
		limit="${CONFIG_GROUP_LIMITS[i]:-$DEFAULT_MAX_PARALLEL}"
		GROUP_NAMES+=("${CONFIG_GROUP_NAMES[i]}")
		GROUP_LIMITS+=("$limit")
		GROUP_STACKS+=("${CONFIG_GROUP_STACKS[i]}")

		read -r -a stacks <<<"${CONFIG_GROUP_STACKS[i]}"
		MANAGED_STACKS+=("${stacks[@]}")
	done

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

	if ((${#UNASSIGNED_STACKS[@]} > 0)); then
		GROUP_NAMES+=("unassigned")
		GROUP_LIMITS+=("$DEFAULT_MAX_PARALLEL")
		GROUP_STACKS+=("${UNASSIGNED_STACKS[*]}")
	fi
}

warn_unassigned() {
	((${#UNASSIGNED_STACKS[@]} == 0)) && return 0

	echo "WARNING: unassigned stacks will run in the final implicit group: ${UNASSIGNED_STACKS[*]}" >&2
}

stack_group() {
	local target="$1"
	local i
	local -a stacks=()

	for ((i = 0; i < ${#GROUP_NAMES[@]}; i++)); do
		read -r -a stacks <<<"${GROUP_STACKS[i]}"
		if contains "$target" "${stacks[@]}"; then
			echo "${GROUP_NAMES[i]}"
			return 0
		fi
	done

	return 1
}

compose() {
	local stack="$1"
	shift
	local -a docker_args=(compose)

	if [[ "${STACKS_QUIET_PROGRESS:-0}" == "1" ]]; then
		docker_args+=(--progress quiet)
	fi

	(
		cd -- "$STACKS_DIR/$stack" || exit 1
		docker "${docker_args[@]}" "$@"
	)
}

run_stack() {
	local stack="$1"
	local action="$2"
	local show_heading="${3:-1}"

	if [[ "$show_heading" == "1" ]]; then
		echo
		echo "==> $stack: $action"
	fi

	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_group() {
	local group_index="$1"
	local action="$2"
	local group_name="${GROUP_NAMES[group_index]}"
	local limit="${GROUP_LIMITS[group_index]}"
	local stack pid finished_pid status result result_color success_result details
	local next=0 running=0 completed=0 failed=0
	local success_color="" failure_color="" color_reset=""
	local -a stacks=()
	local -A pid_stacks=()

	if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
		success_color=$'\033[32m'
		failure_color=$'\033[31m'
		color_reset=$'\033[0m'
	fi

	read -r -a stacks <<<"${GROUP_STACKS[group_index]}"
	((${#stacks[@]} > 0)) || return 0

	if ((limit == 0 || limit > ${#stacks[@]})); then
		limit="${#stacks[@]}"
	fi

	if ((${#stacks[@]} == 1)); then
		details="1 stack"
	elif ((limit == ${#stacks[@]})); then
		details="${#stacks[@]} stacks in parallel"
	else
		details="${#stacks[@]} stacks, max $limit in parallel"
	fi

	echo
	echo "==> Group $group_name: $action ($details)"

	case "$action" in
	up | start) success_result="Started" ;;
	restart) success_result="Restarted" ;;
	stop) success_result="Stopped" ;;
	down) success_result="Removed" ;;
	pull) success_result="Pulled" ;;
	update) success_result="Updated" ;;
	esac

	while ((next < ${#stacks[@]} || running > 0)); do
		while ((next < ${#stacks[@]} && running < limit)); do
			stack="${stacks[next]}"
			STACKS_QUIET_PROGRESS=1 run_stack "$stack" "$action" 0 &
			pid="$!"
			pid_stacks["$pid"]="$stack"
			((next += 1, running += 1))
		done

		finished_pid=""
		status=0
		wait -n -p finished_pid || status=$?
		[[ -n "$finished_pid" ]] || die "cannot identify completed stack"

		stack="${pid_stacks[$finished_pid]}"
		unset 'pid_stacks[$finished_pid]'
		((running -= 1))

		if ((status == 0)); then
			result="$success_result"
			result_color="$success_color"
			((RUN_SUCCESS += 1))
		else
			result="Failed"
			result_color="$failure_color"
			((RUN_FAILED += 1))
			RUN_FAILED_STACKS+=("$stack")
			((failed += 1))
		fi

		((completed += 1))
		printf '[%d/%d] %s: %s%s%s\n' \
			"$completed" "${#stacks[@]}" "$stack" \
			"$result_color" "$result" "$color_reset"
	done

	((failed == 0))
}

print_run_summary() {
	echo
	echo "Result: $RUN_SUCCESS succeeded, $RUN_FAILED failed"

	if ((${#RUN_FAILED_STACKS[@]} > 0)); then
		echo "Failed stacks: ${RUN_FAILED_STACKS[*]}"
	fi

	if ((${#RUN_SKIPPED_GROUPS[@]} > 0)); then
		echo "Skipped groups: ${RUN_SKIPPED_GROUPS[*]}"
	fi
}

collect_skipped_groups() {
	local failed_index="$1"
	local direction="$2"
	local i

	RUN_SKIPPED_GROUPS=()

	if [[ "$direction" == "reverse" ]]; then
		for ((i = failed_index - 1; i >= 0; i--)); do
			RUN_SKIPPED_GROUPS+=("${GROUP_NAMES[i]}")
		done
	else
		for ((i = failed_index + 1; i < ${#GROUP_NAMES[@]}; i++)); do
			RUN_SKIPPED_GROUPS+=("${GROUP_NAMES[i]}")
		done
	fi
}

run_all_grouped() {
	local action="$1"
	local direction="forward"
	local i

	RUN_SUCCESS=0
	RUN_FAILED=0
	RUN_FAILED_STACKS=()
	RUN_SKIPPED_GROUPS=()
	warn_unassigned

	if [[ "$action" == "stop" || "$action" == "down" ]]; then
		direction="reverse"
		for ((i = ${#GROUP_NAMES[@]} - 1; i >= 0; i--)); do
			if ! run_group "$i" "$action"; then
				collect_skipped_groups "$i" "$direction"
				print_run_summary
				return 1
			fi
		done
	else
		for ((i = 0; i < ${#GROUP_NAMES[@]}; i++)); do
			if ! run_group "$i" "$action"; then
				collect_skipped_groups "$i" "$direction"
				print_run_summary
				return 1
			fi
		done
	fi

	print_run_summary
}

run_all_sequential() {
	local action="$1"
	local stack
	local -a failed_stacks=()
	warn_unassigned

	for stack in "${MANAGED_STACKS[@]}"; do
		run_stack "$stack" "$action" || failed_stacks+=("$stack")
	done

	if ((${#failed_stacks[@]} > 0)); then
		echo
		echo "Failed stacks: ${failed_stacks[*]}" >&2
		return 1
	fi
}

list_stacks() {
	local stack group_name status

	warn_unassigned
	printf '%-24s %-16s %s\n' "STACK" "GROUP" "STATUS"
	printf '%-24s %-16s %s\n' "------------------------" "----------------" "---------------"

	for stack in "${DISCOVERED_STACKS[@]}"; do
		if contains "$stack" "${EXCLUDED_STACKS[@]}"; then
			group_name="-"
			status="excluded"
		else
			group_name="$(stack_group "$stack")"
			status="managed"
		fi

		printf '%-24s %-16s %s\n' "$stack" "$group_name" "$status"
	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

	warn_unassigned
	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"
}

load_config
discover_stacks
validate_config_stacks
build_execution_groups

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_grouped "$ACTION"
	fi
	;;
ps)
	if [[ -n "$TARGET" ]]; then
		require_target "$ACTION" "$TARGET"
		run_stack "$TARGET" "$ACTION"
	else
		run_all_sequential "$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