#!/usr/bin/env bash
set -u

# Evacuate non-HA guests after the local Proxmox node enters HA maintenance.
# Guests using shared storage are migrated to one operator-selected node.
# Guests with local storage remain on this node and are gracefully shut down.

HA_WAIT_SECONDS=300
MAINTENANCE_WAIT_SECONDS=60
MAX_PARALLEL_MIGRATIONS=3
SHUTDOWN_TIMEOUT=180
LOCAL_NODE="$(hostname -s)"
MIGRATION_LOG_DIR="$(mktemp -d)"

trap 'rm -rf "$MIGRATION_LOG_DIR"' EXIT

log() {
	printf '\n[%s] %s\n' "$(date '+%F %T')" "$*"
}

warn() {
	printf '\nWARNING: %s\n' "$*" >&2
}

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

command -v pvesh >/dev/null 2>&1 || die "pvesh is required."
command -v ha-manager >/dev/null 2>&1 || die "ha-manager is required."
command -v python3 >/dev/null 2>&1 || die "python3 is required."

pvesh get /cluster/resources --type vm --output-format json >/dev/null 2>&1 ||
	die "Could not read cluster guest resources."
pvesh get /cluster/ha/resources --output-format json >/dev/null 2>&1 ||
	die "Could not read HA resources."
pvesh get /nodes --output-format json >/dev/null 2>&1 ||
	die "Could not read cluster node status."
pvesh get /storage --output-format json >/dev/null 2>&1 ||
	die "Could not read cluster storage configuration."

get_local_guests() {
	pvesh get /cluster/resources --type vm --output-format json 2>/dev/null |
		python3 -c '
import json
import sys

node = sys.argv[1]
for guest in json.load(sys.stdin):
    if guest.get("node") != node or guest.get("type") not in ("qemu", "lxc"):
        continue
    print(
        guest.get("vmid", ""),
        guest.get("type", ""),
        guest.get("status", "unknown"),
        str(guest.get("name", "")).replace("\x1f", " "),
        sep="\x1f",
    )
' "$LOCAL_NODE"
}

get_ha_guest_ids() {
	pvesh get /cluster/ha/resources --output-format json 2>/dev/null |
		python3 -c '
import json
import re
import sys

for resource in json.load(sys.stdin):
    sid = str(resource.get("sid") or resource.get("service") or "")
    match = re.fullmatch(r"(?:vm|ct):(\d+)", sid)
    if match:
        print(match.group(1))
'
}

get_online_nodes() {
	pvesh get /nodes --output-format json 2>/dev/null |
		python3 -c '
import json
import sys

local_node = sys.argv[1]
for node in json.load(sys.stdin):
    name = str(node.get("node", ""))
    if name and name != local_node and node.get("status") == "online":
        print(name)
' "$LOCAL_NODE"
}

get_shared_storages() {
	pvesh get /storage --output-format json 2>/dev/null |
		python3 -c '
import json
import sys

for storage in json.load(sys.stdin):
    if storage.get("shared"):
        print(storage.get("storage", ""))
'
}

guest_storage_scope() {
	local guest_type="$1"
	local vmid="$2"
	local shared_csv="$3"

	pvesh get "/nodes/${LOCAL_NODE}/${guest_type}/${vmid}/config" \
		--output-format json 2>/dev/null |
		python3 -c '
import json
import re
import sys

guest_type, shared_csv = sys.argv[1:3]
shared = {item for item in shared_csv.split(",") if item}
config = json.load(sys.stdin)

if guest_type == "qemu":
    disk_key = re.compile(
        r"^(?:ide|sata|scsi|virtio|efidisk|tpmstate|unused)\d+$"
    )
else:
    disk_key = re.compile(r"^(?:rootfs|mp\d+|unused\d+)$")

local_reasons = []
for key, raw_value in config.items():
    if not disk_key.match(key) or not isinstance(raw_value, str):
        continue

    volume = raw_value.split(",", 1)[0]
    if volume in ("none", "cdrom") or volume.startswith("none,"):
        continue
    if volume.startswith("/") or ":" not in volume:
        local_reasons.append(f"{key}={volume}")
        continue

    storage = volume.split(":", 1)[0]
    if storage not in shared:
        local_reasons.append(f"{key}={storage}")

if local_reasons:
    print("local\x1f" + ", ".join(local_reasons))
else:
    print("shared\x1f")
' "$guest_type" "$shared_csv"
}

select_target_node() {
	local -a nodes
	local choice
	local index

	mapfile -t nodes < <(get_online_nodes)
	(( ${#nodes[@]} > 0 )) || die "No other online cluster node is available."

	printf '\nAvailable migration targets:\n\n'
	for index in "${!nodes[@]}"; do
		printf '  %d) %s\n' "$((index + 1))" "${nodes[$index]}"
	done

	while true; do
		printf '\n'
		read -r -p "Select migration target [1-${#nodes[@]}]: " choice
		if [[ "$choice" =~ ^[0-9]+$ ]] &&
			(( choice >= 1 && choice <= ${#nodes[@]} )); then
			TARGET_NODE="${nodes[$((choice - 1))]}"
			return
		fi
		printf 'Invalid selection.\n' >&2
	done
}

wait_for_ha_evacuation() {
	local deadline=$((SECONDS + HA_WAIT_SECONDS))
	local -a local_guests
	local -a ha_ids
	local -a remaining
	local guest
	local ha_id

	log "Waiting for HA-managed guests to leave ${LOCAL_NODE}."

	while true; do
		mapfile -t local_guests < <(get_local_guests)
		mapfile -t ha_ids < <(get_ha_guest_ids)
		remaining=()

		for guest in "${local_guests[@]}"; do
			for ha_id in "${ha_ids[@]}"; do
				if [[ "${guest%%$'\x1f'*}" == "$ha_id" ]]; then
					remaining+=("$guest")
					break
				fi
			done
		done

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

		if (( SECONDS >= deadline )); then
			warn "HA evacuation did not finish within ${HA_WAIT_SECONDS} seconds."
			printf 'HA-managed guests still assigned to %s:\n' "$LOCAL_NODE" >&2
			printf '  %s\n' "${remaining[@]}" >&2
			return 1
		fi

		printf '\r  Waiting: %d HA guest(s) remain... ' "${#remaining[@]}"
		sleep 5
	done
}

migrate_guest() {
	local vmid="$1"
	local guest_type="$2"
	local status="$3"

	if [[ "$guest_type" == "qemu" ]]; then
		if [[ "$status" == "running" ]]; then
			qm migrate "$vmid" "$TARGET_NODE" --online
		else
			qm migrate "$vmid" "$TARGET_NODE"
		fi
	else
		if [[ "$status" == "running" ]]; then
			pct migrate "$vmid" "$TARGET_NODE" --restart 1 \
				--timeout "$SHUTDOWN_TIMEOUT"
		else
			pct migrate "$vmid" "$TARGET_NODE"
		fi
	fi
}

wait_for_migration_batch() {
	local index
	local pid
	local guest
	local log_file
	local vmid
	local guest_type
	local status
	local name

	for index in "${!batch_pids[@]}"; do
		pid="${batch_pids[$index]}"
		guest="${batch_guests[$index]}"
		log_file="${batch_logs[$index]}"
		IFS=$'\x1f' read -r vmid guest_type status name <<< "$guest"

		if wait "$pid"; then
			log "Migration completed for ${guest_type} ${vmid} (${name:-unnamed})."
		else
			warn "Migration failed for ${guest_type} ${vmid} (${name:-unnamed})."
			migration_failures+=("$guest")
		fi

		if [[ -s "$log_file" ]]; then
			sed 's/^/  /' "$log_file"
		fi
	done

	batch_pids=()
	batch_guests=()
	batch_logs=()
}

shutdown_guest() {
	local vmid="$1"
	local guest_type="$2"

	if [[ "$guest_type" == "qemu" ]]; then
		qm shutdown "$vmid" --timeout "$SHUTDOWN_TIMEOUT"
	else
		pct shutdown "$vmid" --timeout "$SHUTDOWN_TIMEOUT"
	fi
}

guest_status() {
	local vmid="$1"
	local guest_type="$2"

	if [[ "$guest_type" == "qemu" ]]; then
		qm status "$vmid" 2>/dev/null | awk '{ print $2 }'
	else
		pct status "$vmid" 2>/dev/null | awk '{ print $2 }'
	fi
}

wait_for_maintenance_mode() {
	local deadline=$((SECONDS + MAINTENANCE_WAIT_SECONDS))

	log "Waiting for ${LOCAL_NODE} to enter HA maintenance mode."

	while true; do
		if ha-manager status |
			grep -F "$LOCAL_NODE" |
			grep -q "maintenance mode"; then
			printf '\n'
			return
		fi

		if (( SECONDS >= deadline )); then
			die "${LOCAL_NODE} did not enter HA maintenance mode within ${MAINTENANCE_WAIT_SECONDS} seconds."
		fi

		printf '\r  Waiting for HA maintenance mode... '
		sleep 2
	done
}

wait_for_maintenance_mode

wait_for_ha_evacuation ||
	die "Resolve the remaining HA guests before continuing the evacuation."

select_target_node

mapfile -t shared_storages < <(get_shared_storages)
shared_csv="$(IFS=,; echo "${shared_storages[*]}")"

mapfile -t guests < <(get_local_guests)
if (( ${#guests[@]} == 0 )); then
	log "No guests remain on ${LOCAL_NODE}."
	exit 0
fi

declare -a shared_guests=()
declare -a local_guests=()

for guest in "${guests[@]}"; do
	IFS=$'\x1f' read -r vmid guest_type status name <<< "$guest"
	storage_result="$(guest_storage_scope "$guest_type" "$vmid" "$shared_csv")" ||
		die "Could not inspect storage for ${guest_type} ${vmid}."
	IFS=$'\x1f' read -r scope reason <<< "$storage_result"

	if [[ "$scope" == "shared" ]]; then
		shared_guests+=("$guest")
	else
		local_guests+=("${guest}"$'\x1f'"${reason:-local storage}")
	fi
done

printf '\nEvacuation plan for %s:\n' "$LOCAL_NODE"
printf '  Migration target: %s\n' "$TARGET_NODE"
printf '  Shared-storage guests to migrate: %d\n' "${#shared_guests[@]}"
printf '  Local-storage guests to retain:   %d\n' "${#local_guests[@]}"

if (( ${#shared_guests[@]} > 0 )); then
	printf '\nShared-storage guests:\n'
	for guest in "${shared_guests[@]}"; do
		IFS=$'\x1f' read -r vmid guest_type status name <<< "$guest"
		printf '  %-6s %-5s %-8s %s\n' "$vmid" "$guest_type" "$status" "$name"
	done
fi

if (( ${#local_guests[@]} > 0 )); then
	printf '\nLocal-storage guests (will not migrate):\n'
	for guest in "${local_guests[@]}"; do
		IFS=$'\x1f' read -r vmid guest_type status name reason <<< "$guest"
		printf '  %-6s %-5s %-8s %-24s %s\n' \
			"$vmid" "$guest_type" "$status" "$name" "$reason"
	done
fi

printf '\n'
read -r -p "Proceed with shared-storage guest migration? [y/N] " answer
[[ "$answer" =~ ^[Yy]$ ]] || { echo "Evacuation cancelled."; exit 0; }

declare -a migration_failures=()
declare -a batch_pids=()
declare -a batch_guests=()
declare -a batch_logs=()

for guest in "${shared_guests[@]}"; do
	IFS=$'\x1f' read -r vmid guest_type status name <<< "$guest"
	log "Starting migration for ${guest_type} ${vmid} (${name:-unnamed}) to ${TARGET_NODE}."

	log_file="${MIGRATION_LOG_DIR}/${guest_type}-${vmid}.log"
	migrate_guest "$vmid" "$guest_type" "$status" >"$log_file" 2>&1 &
	batch_pids+=("$!")
	batch_guests+=("$guest")
	batch_logs+=("$log_file")

	if (( ${#batch_pids[@]} >= MAX_PARALLEL_MIGRATIONS )); then
		wait_for_migration_batch
	fi
done

(( ${#batch_pids[@]} == 0 )) || wait_for_migration_batch

declare -a running_local_guests=()
for guest in "${local_guests[@]}"; do
	IFS=$'\x1f' read -r vmid guest_type status name reason <<< "$guest"
	status="$(guest_status "$vmid" "$guest_type")"
	if [[ "$status" == "running" ]]; then
		running_local_guests+=(
			"${vmid}"$'\x1f'"${guest_type}"$'\x1f'"${status}"$'\x1f'"${name}"$'\x1f'"${reason}"
		)
	fi
done

if (( ${#running_local_guests[@]} > 0 )); then
	printf '\nThe following local-storage guests remain running:\n'
	for guest in "${running_local_guests[@]}"; do
		IFS=$'\x1f' read -r vmid guest_type status name reason <<< "$guest"
		printf '  %-6s %-5s %-24s %s\n' "$vmid" "$guest_type" "$name" "$reason"
	done

	printf '\n'
	read -r -p "Gracefully shut down these local-storage guests? [y/N] " answer
	if [[ "$answer" =~ ^[Yy]$ ]]; then
		for guest in "${running_local_guests[@]}"; do
			IFS=$'\x1f' read -r vmid guest_type status name reason <<< "$guest"
			log "Shutting down ${guest_type} ${vmid} (${name:-unnamed})."
			shutdown_guest "$vmid" "$guest_type" ||
				warn "Graceful shutdown failed for ${guest_type} ${vmid}; it was not force-stopped."
		done
	else
		warn "Local-storage guests were left running."
	fi
fi

mapfile -t final_guests < <(get_local_guests)

printf '\nEvacuation summary:\n'
printf '  Migration failures: %d\n' "${#migration_failures[@]}"
printf '  Guests still assigned to %s: %d\n' "$LOCAL_NODE" "${#final_guests[@]}"

if (( ${#migration_failures[@]} > 0 )); then
	printf '\nFailed migrations (left unchanged; not shut down automatically):\n'
	printf '  %s\n' "${migration_failures[@]}"
fi

if (( ${#final_guests[@]} > 0 )); then
	printf '\nGuests still assigned to %s:\n' "$LOCAL_NODE"
	printf '  %s\n' "${final_guests[@]}"
fi

printf '\nNo guest was force-stopped.\n'