#!/usr/bin/env bash

set -Eeuo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
SCRIPT_PATH="${SCRIPT_DIR}/$(basename -- "${BASH_SOURCE[0]}")"

fail() {
	printf 'ERROR: %s\n' "$*" >&2
	exit 1
}

trap 'printf "ERROR: broker update failed at line %d.\n" "$LINENO" >&2' ERR

cd "$SCRIPT_DIR"

command -v git >/dev/null 2>&1 || fail "git is required"

POST_PULL=0
LOCAL_ONLY="${TAPM_UPDATE_LOCAL_ONLY:-0}"
for argument in "$@"; do
	case "$argument" in
		--post-pull) POST_PULL=1 ;;
		--local-only) LOCAL_ONLY=1 ;;
		*) fail "unknown argument: ${argument}" ;;
	esac
done

deploy_local_broker() {
	command -v docker >/dev/null 2>&1 || fail "docker is required"
	docker compose version >/dev/null 2>&1 || fail "Docker Compose v2 is required"
	[[ -f compose.yaml ]] || fail "compose.yaml was not found in ${SCRIPT_DIR}"
	[[ -f .env ]] || fail ".env was not found in ${SCRIPT_DIR}"

	printf 'Building broker image...\n'
	docker compose build broker

	printf 'Applying pending database migrations...\n'
	docker compose run --rm migrate

	printf 'Starting broker...\n'
	docker compose up -d broker

	container_id="$(docker compose ps -q broker)"
	[[ -n "$container_id" ]] || fail "Docker Compose did not return a broker container"

	printf 'Waiting for broker health check'
	for ((attempt = 1; attempt <= 75; attempt++)); do
		status="$(docker inspect --format '{{if .State.Health}}{{.State.Health.Status}}{{else}}{{.State.Status}}{{end}}' "$container_id")"
		case "$status" in
			healthy)
				printf ' healthy.\n'
				docker compose ps broker
				printf 'Local TAPM Deployment Broker update completed successfully.\n'
				return 0
				;;
			exited | dead | unhealthy)
				printf ' %s.\n' "$status"
				docker compose logs --tail=100 broker
				fail "broker failed its health check"
				;;
		esac
		printf '.'
		sleep 2
	done

	printf ' timed out.\n'
	docker compose logs --tail=100 broker
	fail "broker did not become healthy within 150 seconds"
}

update_peer_broker() {
	local listen_addr
	local local_address
	local peer_host="${TAPM_UPDATE_PEER_HOST:-}"
	local ssh_user="${TAPM_UPDATE_SSH_USER:-root}"
	local remote_update_path='/opt/idssys/ta-deployment-broker/update.sh'
	local -a ssh_options=(
		-o BatchMode=yes
		-o ConnectTimeout=5
		-o ConnectionAttempts=1
		-o StrictHostKeyChecking=yes
	)

	[[ -f .env ]] || fail ".env was not found in ${SCRIPT_DIR}"
	if [[ -z "$peer_host" ]]; then
		listen_addr="$(
			sed -n 's/^[[:space:]]*TAPM_LISTEN_ADDR[[:space:]]*=[[:space:]]*//p' .env |
				tail -n 1
		)"
		listen_addr="${listen_addr#\"}"
		listen_addr="${listen_addr%\"}"
		listen_addr="${listen_addr#\'}"
		listen_addr="${listen_addr%\'}"
		local_address="${listen_addr%:*}"
		case "$local_address" in
			10.10.1.121) peer_host='10.10.1.122' ;;
			10.10.1.122) peer_host='10.10.1.121' ;;
			*) fail "unable to determine peer from TAPM_LISTEN_ADDR=${listen_addr}" ;;
		esac
	fi

	printf 'Checking peer broker at %s...\n' "$peer_host"
	if ! ssh "${ssh_options[@]}" "${ssh_user}@${peer_host}" true; then
		printf 'WARNING: peer %s is unavailable over SSH; its update was skipped.\n' "$peer_host" >&2
		return 0
	fi
	if ! ssh "${ssh_options[@]}" "${ssh_user}@${peer_host}" \
		"TAPM_UPDATE_LOCAL_ONLY=1 ${remote_update_path} --local-only"; then
		fail "peer ${peer_host} is reachable, but its broker update failed"
	fi
	printf 'Peer broker at %s completed successfully.\n' "$peer_host"
}

if ((POST_PULL == 0)); then
	branch="$(git symbolic-ref --quiet --short HEAD)" ||
		fail "the repository is in detached HEAD state"
	upstream="$(git rev-parse --abbrev-ref --symbolic-full-name '@{upstream}' 2>/dev/null)" ||
		fail "branch ${branch} does not have an upstream branch"
	remote="$(git config --get "branch.${branch}.remote")" ||
		fail "branch ${branch} does not have a configured remote"
	git diff --quiet && git diff --cached --quiet ||
		fail "tracked local changes must be committed or stashed before updating"

	printf 'Checking %s against %s...\n' "$branch" "$upstream"
	git fetch --prune "$remote"

	local_commit="$(git rev-parse HEAD)"
	upstream_commit="$(git rev-parse '@{upstream}')"
	if [[ "$local_commit" == "$upstream_commit" ]]; then
		printf 'Local TAPM Deployment Broker is already up to date.\n'
	elif git merge-base --is-ancestor "$local_commit" "$upstream_commit"; then
		printf 'Updating TAPM Deployment Broker from %s...\n' "$upstream"
		git merge --ff-only "$upstream_commit"

		# Restart from the newly pulled script so updates to this file take effect.
		if ((LOCAL_ONLY == 1)); then
			exec "$SCRIPT_PATH" --post-pull --local-only
		fi
		exec "$SCRIPT_PATH" --post-pull
	elif git merge-base --is-ancestor "$upstream_commit" "$local_commit"; then
		fail "local branch ${branch} is ahead of ${upstream}; refusing to rebuild"
	else
		fail "local branch ${branch} has diverged from ${upstream}"
	fi
else
	deploy_local_broker
fi

if ((LOCAL_ONLY == 0)); then
	update_peer_broker
fi

printf 'TAPM Deployment Broker update check completed.\n'