#!/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=()

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

# manage.conf DSL: 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"
}

# manage.conf DSL: 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+=("$*")
}

# manage.conf DSL: 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

	bash -n "$CONFIG_FILE" || die "syntax error in config file: $CONFIG_FILE"
	# This is a trusted Bash configuration file that calls the DSL above.
	source "$CONFIG_FILE"

	local legacy
	for legacy in WHITELIST BLACKLIST PRIORITY; do
		if declare -p "$legacy" >/dev/null 2>&1; then
			die "$legacy is no longer supported; migrate $CONFIG_FILE to group/exclude directives"
		fi
	done
}

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_FORCE_ANSI:-0}" == "1" ]]; then
		docker_args+=(--ansi always)
	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_stack_prefixed() {
	local stack="$1"
	local action="$2"
	local line status

	STACKS_FORCE_ANSI=1 run_stack "$stack" "$action" 0 2>&1 |
		while IFS= read -r line || [[ -n "$line" ]]; do
			[[ -n "$line" ]] || continue
			printf '[%s] %s\n' "$stack" "$line"
		done
	status="${PIPESTATUS[0]}"

	return "$status"
}

run_group() {
	local group_index="$1"
	local action="$2"
	local group_name="${GROUP_NAMES[group_index]}"
	local limit="${GROUP_LIMITS[group_index]}"
	local semaphore_dir semaphore_fifo semaphore_fd
	local stack pid i status
	local failed=0
	local -a stacks=() pids=() pid_stacks=()

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

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

	echo
	echo "==> Group $group_name: $action (${#stacks[@]} stacks, parallel=$limit)"

	semaphore_dir="$(mktemp -d "${TMPDIR:-/tmp}/stacks-manage.XXXXXX")" ||
		die "cannot create semaphore directory"
	semaphore_fifo="$semaphore_dir/tokens"
	mkfifo "$semaphore_fifo" || die "cannot create semaphore"
	exec {semaphore_fd}<>"$semaphore_fifo" || die "cannot open semaphore"

	for ((i = 0; i < limit; i++)); do
		printf '\n' >&"$semaphore_fd"
	done

	for stack in "${stacks[@]}"; do
		IFS= read -r -u "$semaphore_fd"

		(
			status=0
			run_stack_prefixed "$stack" "$action" || status=$?
			printf '\n' >&"$semaphore_fd"
			exit "$status"
		) &

		pids+=("$!")
		pid_stacks+=("$stack")
	done

	for ((i = 0; i < ${#pids[@]}; i++)); do
		pid="${pids[i]}"
		stack="${pid_stacks[i]}"

		if wait "$pid"; then
			((RUN_SUCCESS += 1))
		else
			((RUN_FAILED += 1))
			RUN_FAILED_STACKS+=("$stack")
			((failed += 1))
		fi
	done

	exec {semaphore_fd}>&-
	rm -f -- "$semaphore_fifo"
	rmdir -- "$semaphore_dir"

	((failed == 0))
}

print_run_summary() {
	echo
	echo "Succeeded: $RUN_SUCCESS"
	echo "Failed:    $RUN_FAILED"

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

print_skipped_groups() {
	local failed_index="$1"
	local direction="$2"
	local i
	local -a skipped=()

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

	if ((${#skipped[@]} > 0)); then
		echo "Skipped groups after failure: ${skipped[*]}" >&2
	fi
}

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

	RUN_SUCCESS=0
	RUN_FAILED=0
	RUN_FAILED_STACKS=()
	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
				print_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
				print_skipped_groups "$i" "$direction"
				print_run_summary
				return 1
			fi
		done
	fi

	print_run_summary
}

run_all_sequential() {
	local action="$1"
	local stack

	RUN_SUCCESS=0
	RUN_FAILED=0
	RUN_FAILED_STACKS=()
	warn_unassigned

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

	print_run_summary
	((RUN_FAILED == 0))
}

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