#!/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
WAIT_TIMEOUT=120

BEFORE_STAGES=()
BEFORE_LIMITS=()
AFTER_STAGES=()
AFTER_LIMITS=()
CONFIGURED_STACKS=()
EXCLUDED_STACKS=()

DISCOVERED_STACKS=()
MANAGED_STACKS=()
STAGE_NAMES=()
STAGE_LIMITS=()
STAGE_STACKS=()

RUN_SUCCESS=0
RUN_FAILED=0
RUN_FAILED_STACKS=()
RUN_SKIPPED_STAGES=()

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

if ((BASH_VERSINFO[0] < 5 || (\
	BASH_VERSINFO[0] == 5 && BASH_VERSINFO[1] < 1))); then
	die "Bash 5.1 or newer is required"
fi

handle_signal() {
	local status="$1"
	local pid attempt alive
	local -a pids=()

	trap '' INT TERM
	mapfile -t pids < <(jobs -p)
	for pid in "${pids[@]}"; do
		kill -TERM -- "-$pid" 2>/dev/null || :
	done

	# Allow five seconds for cleanup, then stop every remaining descendant.
	for ((attempt = 0; attempt < 50; attempt++)); do
		alive=0
		for pid in "${pids[@]}"; do
			kill -0 -- "-$pid" 2>/dev/null && alive=1
		done
		((alive)) || break
		sleep 0.1
	done
	for pid in "${pids[@]}"; do
		kill -KILL -- "-$pid" 2>/dev/null || :
	done
	((${#pids[@]} == 0)) || wait "${pids[@]}" 2>/dev/null
	printf '\n'
	exit "$status"
}

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

Lifecycle:
  up [STACK]       Create and start containers
  start [STACK]    Start existing containers
  stop [STACK]     Stop containers
  restart [STACK]  Restart containers
  down [STACK]     Stop and remove containers and networks
  pull [STACK]     Pull images
  update [STACK]   Pull images and apply Compose changes

Inspect:
  ps [STACK]       Show container status
  logs STACK       Follow logs
  config STACK     Render Compose configuration
  images [STACK]   Show services, configured images, and versions
  list             List stacks, stages, and management status
  help, -h, --help Show this help

Stage configuration:
  max_parallel N                   Set concurrency, 0-65535 (0 = unlimited)
  wait_timeout N                   Wait 0-86400 seconds for startup (default 120;
                                   0 disables waiting)
  before [--parallel N] STACK...   Add an early stage
  after [--parallel N] STACK...    Add a late stage
  exclude STACK...                 Exclude stacks from bulk operations

Bulk execution (without STACK):
  Each before/after declaration creates a stage, in declaration order.
  Unlisted stacks form the default stage; excluded stacks are skipped.
  up/start/restart/pull/update run before -> default -> after.
  stop/down reverse the full stage order.
  Stages run one at a time. Stacks within a stage run in parallel, limited
  by max_parallel or that stage's --parallel setting.
  A stage finishes all its stacks; any failure skips the remaining stages.

Startup checks:
  up/update validate all selected Compose configurations before making changes.
  up/start/update wait up to wait_timeout seconds for each stack to be ready.
  Services with healthchecks must be healthy; other services must be running.
  Set wait_timeout to 0 to disable waiting.

Config: $CONFIG_FILE
EOF
}

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

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

	return 1
}

validate_integer_range() {
	local kind="$1"
	local value="$2"
	local maximum="$3"

	[[ "$value" =~ ^(0|[1-9][0-9]*)$ ]] &&
		((${#value} <= ${#maximum})) && ((value <= maximum)) ||
		die "$kind must be an integer between 0 and $maximum without leading zeros: $value"
}

validate_parallel_limit() {
	validate_integer_range "parallel limit" "$1" 65535
}

wait_timeout() {
	[[ $# -eq 1 ]] || die "wait_timeout expects exactly one value"
	validate_integer_range "wait timeout" "$1" 86400
	WAIT_TIMEOUT="$1"
}

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

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

# Set the maximum concurrency within each stage.
max_parallel() {
	[[ $# -eq 1 ]] || die "max_parallel expects exactly one value"
	validate_parallel_limit "$1"
	DEFAULT_MAX_PARALLEL="$1"
}

add_stage() {
	local stages_name="$1"
	local limits_name="$2"
	local stage="$3"
	local limit=""
	local stack
	local -n stages="$stages_name"
	local -n limits="$limits_name"
	shift 3

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

	[[ $# -ge 1 ]] || die "$stage expects at least one stack"

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

	stages+=("$*")
	limits+=("$limit")
}

before() { add_stage BEFORE_STAGES BEFORE_LIMITS before "$@"; }
after() { add_stage AFTER_STAGES AFTER_LIMITS after "$@"; }

# 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 assigned and excluded: $stack"
		EXCLUDED_STACKS+=("$stack")
	done
}

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

	# The trusted config calls the configuration functions 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}"
		stack="${stack##*/}"
		validate_name "stack directory" "$stack"
		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
}

append_stages() {
	local stage="$1"
	local stages_name="$2"
	local limits_name="$3"
	local i name
	local -a stacks=()
	local -n stages="$stages_name"
	local -n limits="$limits_name"

	for ((i = 0; i < ${#stages[@]}; i++)); do
		name="$stage"
		((${#stages[@]} > 1)) && name+="-$((i + 1))"
		STAGE_NAMES+=("$name")
		STAGE_LIMITS+=("${limits[i]:-$DEFAULT_MAX_PARALLEL}")
		STAGE_STACKS+=("${stages[i]}")
		read -r -a stacks <<<"${stages[i]}"
		MANAGED_STACKS+=("${stacks[@]}")
	done
}

build_execution_stages() {
	local stack
	local -a default_stacks=()

	STAGE_NAMES=()
	STAGE_LIMITS=()
	STAGE_STACKS=()
	MANAGED_STACKS=()

	append_stages before BEFORE_STAGES BEFORE_LIMITS

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

	if ((${#default_stacks[@]} > 0)); then
		STAGE_NAMES+=("default")
		STAGE_LIMITS+=("$DEFAULT_MAX_PARALLEL")
		STAGE_STACKS+=("${default_stacks[*]}")
		MANAGED_STACKS+=("${default_stacks[@]}")
	fi

	append_stages after AFTER_STAGES AFTER_LIMITS
}

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

	for ((i = 0; i < ${#STAGE_NAMES[@]}; i++)); do
		read -r -a stacks <<<"${STAGE_STACKS[i]}"
		if contains "$target" "${stacks[@]}"; then
			echo "${STAGE_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}"
	local -a wait_args=()

	if ((WAIT_TIMEOUT > 0)); then
		wait_args=(--wait --wait-timeout "$WAIT_TIMEOUT")
	fi

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

	case "$action" in
	up)
		compose "$stack" up -d "${wait_args[@]}"
		;;
	start)
		compose "$stack" start "${wait_args[@]}"
		;;
	update)
		compose "$stack" pull &&
			compose "$stack" up -d "${wait_args[@]}"
		;;
	*)
		compose "$stack" "$action"
		;;
	esac
}

preflight_configs() {
	local stack
	local -a failed_stacks=()

	(($# > 0)) || return 0
	printf '\n==> Checking Compose configurations\n'
	for stack in "$@"; do
		compose "$stack" config -q || failed_stacks+=("$stack")
	done
	if ((${#failed_stacks[@]} > 0)); then
		printf 'ERROR: invalid Compose configuration: %s\n' "${failed_stacks[*]}" >&2
		return 1
	fi
}

run_stage() {
	local stage_index="$1"
	local action="$2"
	local stage_name="${STAGE_NAMES[stage_index]}"
	local limit="${STAGE_LIMITS[stage_index]}"
	local stack pid finished_pid status success_result working_result details
	local index finished_index total
	local next=0 running=0 completed=0 failed=0 stack_width=0 interactive=0 rendered=0
	local success_color="" failure_color="" color_reset=""
	local -a stacks=() states=() markers=() state_colors=()
	local -A pid_indexes=()

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

	read -r -a stacks <<<"${STAGE_STACKS[stage_index]}"
	total="${#stacks[@]}"
	((total > 0)) || return 0
	for ((index = 0; index < total; index++)); do
		stack="${stacks[index]}"
		((${#stack} > stack_width)) && stack_width="${#stack}"
		states+=("Waiting")
		markers+=(" ")
		state_colors+=("")
	done

	if ((limit == 0 || limit > total)); then
		limit="$total"
	fi

	if ((total == 1)); then
		details="1 stack"
	elif ((limit == total)); then
		details="$total stacks in parallel"
	else
		details="$total stacks, max $limit in parallel"
	fi

	echo
	if ((!interactive)); then
		echo "==> Stage $stage_name: $action ($details)"
	fi

	case "$action" in
	up | start)
		working_result="Starting"
		success_result="Started"
		;;
	restart)
		working_result="Restarting"
		success_result="Restarted"
		;;
	stop)
		working_result="Stopping"
		success_result="Stopped"
		;;
	down)
		working_result="Removing"
		success_result="Removed"
		;;
	pull)
		working_result="Pulling"
		success_result="Pulled"
		;;
	update)
		working_result="Updating"
		success_result="Updated"
		;;
	esac

	while :; do
		while ((next < total && running < limit)); do
			stack="${stacks[next]}"
			states[next]="$working_result"
			STACKS_QUIET_PROGRESS=1 run_stack "$stack" "$action" 0 &
			pid="$!"
			pid_indexes["$pid"]="$next"
			((next += 1, running += 1))
		done

		if ((interactive)); then
			((rendered)) && printf '\033[%dA' "$((total + 1))"
			printf '\r\033[2K[+] %s: %s %d/%d\n' \
				"$stage_name" "$action" "$completed" "$total"
			for ((index = 0; index < total; index++)); do
				printf '\r\033[2K %s%s%s Stack %-*s %s%s%s\n' \
					"${state_colors[index]}" "${markers[index]}" "$color_reset" \
					"$stack_width" "${stacks[index]}" \
					"${state_colors[index]}" "${states[index]}" "$color_reset"
			done
			rendered=1
		fi

		((running > 0)) || break

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

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

		if ((status == 0)); then
			markers[finished_index]="✔"
			states[finished_index]="$success_result"
			state_colors[finished_index]="$success_color"
			((RUN_SUCCESS += 1))
		else
			markers[finished_index]="✘"
			states[finished_index]="Failed"
			state_colors[finished_index]="$failure_color"
			((RUN_FAILED += 1))
			RUN_FAILED_STACKS+=("$stack")
			((failed += 1))
		fi

		((completed += 1))
		if ((!interactive)); then
			printf ' %s%s%s Stack %-*s %s%s%s\n' \
				"${state_colors[finished_index]}" "${markers[finished_index]}" "$color_reset" \
				"$stack_width" "$stack" \
				"${state_colors[finished_index]}" "${states[finished_index]}" "$color_reset"
		fi
	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_STAGES[@]} > 0)); then
		echo "Skipped stages: ${RUN_SKIPPED_STAGES[*]}"
	fi
}

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

	RUN_SKIPPED_STAGES=()

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

run_all_staged() {
	local action="$1"
	local direction="forward"
	local start=0 end="${#STAGE_NAMES[@]}" step=1
	local i

	RUN_SUCCESS=0
	RUN_FAILED=0
	RUN_FAILED_STACKS=()
	RUN_SKIPPED_STAGES=()
	if [[ "$action" == "stop" || "$action" == "down" ]]; then
		direction="reverse"
		start=$((end - 1))
		end=-1
		step=-1
	fi

	for ((i = start; i != end; i += step)); do
		if ! run_stage "$i" "$action"; then
			collect_skipped_stages "$i" "$direction"
			print_run_summary
			return 1
		fi
	done

	print_run_summary
}

run_all_sequential() {
	local action="$1"
	local stack
	local -a failed_stacks=()
	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
}

print_table() {
	local columns="$1"
	shift
	local -a cells=("$@") widths=()
	local index column row

	for ((index = 0; index < ${#cells[@]}; index++)); do
		column=$((index % columns))
		if ((${#cells[index]} > ${widths[column]:-0})); then
			widths[column]="${#cells[index]}"
		fi
	done

	for ((row = 0; row < ${#cells[@]}; row += columns)); do
		for ((column = 0; column < columns; column++)); do
			((column > 0)) && printf '   '
			if ((column == columns - 1)); then
				printf '%s\n' "${cells[row + column]}"
			else
				printf '%-*s' "${widths[column]}" "${cells[row + column]}"
			fi
		done
	done
}

list_stacks() {
	local stack stage_name status
	local -a cells=(STACK STAGE STATUS)

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

		cells+=("$stack" "$stage_name" "$status")
	done

	print_table 3 "${cells[@]}"
}

list_images() {
	local target="${1:-}"
	local stack output service image
	local failed=0
	local -a stacks=("${MANAGED_STACKS[@]}") cells=(STACK SERVICE IMAGE)

	if [[ -n "$target" ]]; then
		require_target "images" "$target"
		stacks=("$target")
	fi

	for stack in "${stacks[@]}"; do
		# Read service/image pairs from Compose's canonical YAML (two-space indent).
		if ! output="$(compose "$stack" config --format yaml | awk -v project="$stack" '
			{ gsub(/["\047]/, "") }
			/^name: / { project = $2 }
			/^[^ ]/ { in_services = ($0 == "services:") }
			in_services && /^  [^ ]/ {
				service = $1
				sub(/:$/, "", service)
				services[++count] = service
				images[service] = project "-" service
			}
			in_services && /^    image: / { images[service] = $2 }
			END {
				for (i = 1; i <= count; i++)
					print services[i] "\t" images[services[i]]
			}
		')"; then
			((failed += 1))
			continue
		fi
		while IFS=$'\t' read -r service image; do
			[[ -n "$service" ]] || continue
			cells+=("$stack" "$service" "$image")
		done <<<"$output"
	done

	print_table 3 "${cells[@]}"
	((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"
}

main() {
	local ACTION="${1:-}"
	local TARGET="${2:-}"
	local -a targets=()

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

	case "$ACTION" in
	help | -h | --help)
		usage
		exit 0
		;;
	esac

	load_config
	discover_stacks
	validate_config_stacks
	build_execution_stages

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

		if [[ -n "$TARGET" ]]; then
			require_target "$ACTION" "$TARGET"
			targets=("$TARGET")
		else
			targets=("${MANAGED_STACKS[@]}")
		fi
		if [[ "$ACTION" == "up" || "$ACTION" == "update" ]]; then
			preflight_configs "${targets[@]}" || return 1
		fi

		if [[ -n "$TARGET" ]]; then
			run_stack "$TARGET" "$ACTION"
		else
			run_all_staged "$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
		;;
	*)
		usage
		exit 1
		;;
	esac
}

if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
	trap 'handle_signal 130' INT
	trap 'handle_signal 143' TERM
	# One isolated process group covers the scheduler and all Docker descendants.
	# Disable job control inside it so parallel workers stay in the same group.
	set -m
	(
		set +m
		trap - INT TERM
		main "$@"
	) <&0 &
	main_pid=$!
	set +m
	wait "$main_pid"
fi