#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-2.0-or-later
# Copyright (C) 2026 Bryan Joshua Pedini

# tests/docker/cleanup.sh - restore the Docker image set recorded by snapshot.sh.
#
# Algorithm, in this order:
#   a. read the "before" maps from the snapshot; compute the "after" maps the same way
#   b. re-tag stolen names: every pre-existing repo:tag that now points at another
#      image (or vanished) is pointed back with `docker tag <id_before> <repo:tag>`
#      (e.g. a pre-existing archlinux:latest whose tag was taken by a newer pull)
#   c. untag new names: every repo:tag not in the snapshot -> `docker rmi --no-prune`
#      (only untags when the image has other names; deletes it on its last name)
#   d. delete new ids: every image id not in the snapshot -> `docker rmi -f --no-prune`
#   e. print a summary; exit 1 if a pre-existing repo:tag is missing or points elsewhere
#
# Safety rules:
#   - containers are never stopped or removed; an image used by ANY container
#     (running or stopped) is left in place and reported as left over;
#   - an image id that is in the snapshot is never deleted: --no-prune keeps
#     untagged parents, and the last name of a pre-existing image is never removed;
#   - idempotent: running it again, or when nothing changed, is a no-op.
#
# usage: cleanup.sh [--dry-run]
# exit:  0 ok, 1 a pre-existing repo:tag is still missing/wrong, 2 usage/environment error
set -euo pipefail

prog=${0##*/}
here=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
tags_file=$here/images-before.tsv
ids_file=$here/image-ids-before.txt

usage() {
	cat <<USAGE
usage: $prog [--dry-run]

Restore the Docker image set to the snapshot in
  $tags_file
  $ids_file
(re-tag stolen pre-existing tags, untag new tags, delete new image ids).

  -n, --dry-run   print the docker commands instead of running them; the plan is
                  simulated in memory and the summary shows the predicted end state
  -h, --help      this help
USAGE
}

dry_run=0
for arg in "$@"; do
	case $arg in
	-n | --dry-run) dry_run=1 ;;
	-h | --help) usage; exit 0 ;;
	*) printf '%s: unknown argument: %s\n' "$prog" "$arg" >&2; usage >&2; exit 2 ;;
	esac
done

warn() { printf '%s: WARNING: %s\n' "$prog" "$*" >&2; }
die() { printf '%s: error: %s\n' "$prog" "$*" >&2; exit 2; }
short() { local h=${1#sha256:}; printf '%s' "${h:0:12}"; }

[[ -f $tags_file && -f $ids_file ]] ||
	die "no snapshot (${tags_file##*/}, ${ids_file##*/} in $here); run tests/docker/snapshot.sh first"
command -v docker >/dev/null 2>&1 || die "docker CLI not found"
docker info >/dev/null 2>&1 || die "cannot talk to the Docker daemon"

# ---- state -----------------------------------------------------------------
declare -A before_tag=()   # repo:tag -> id          (snapshot)
declare -A before_id=()    # id -> 1                 (snapshot: every id, incl. untagged)
declare -A tag_id=()       # repo:tag -> id          (current; simulated under --dry-run)
declare -A tag_count=()    # id -> number of repo:tag names on it
declare -A id_present=()   # id -> 1
declare -A used_running=() # id -> "name (status), ..."  running/paused/restarting containers
declare -A used_stopped=() # id -> "name (status), ..."  created/exited/dead containers
declare -A id_names=()     # id -> "repo:tag ..." as found at start (for messages only)
declare -a done_retag=() done_untag=() done_remove=() skipped=() failed=()

# keys_sorted NAME: keys of associative array NAME, one per line, C-sorted.
keys_sorted() {
	local -n _a=$1
	((${#_a[@]})) || return 0
	printf '%s\n' "${!_a[@]}" | LC_ALL=C sort
}

# users_of ID: "running container(s): ...; stopped container(s): ..." or empty.
users_of() {
	local s=
	[[ -z ${used_running[$1]:-} ]] || s="running container(s): ${used_running[$1]}"
	[[ -z ${used_stopped[$1]:-} ]] || s+="${s:+; }stopped container(s): ${used_stopped[$1]}"
	printf '%s' "$s"
}

# run CMD...: print "+ CMD"; unless --dry-run, run it, show its output indented
# and leave a one-line error in $run_err on failure.
run_err=
run() {
	local out rc=0
	local -a lines
	run_err=
	printf '  + %s\n' "$*"
	((!dry_run)) || return 0
	out=$("$@" 2>&1) || rc=$?
	if [[ -n $out ]]; then
		mapfile -t lines <<<"$out"
		printf '      %s\n' "${lines[@]}"
	fi
	((rc == 0)) || run_err=${out//$'\n'/ }
	return "$rc"
}

skip() {
	printf '  - skip: %s\n' "$*"
	skipped+=("$*")
}

read_before() {
	local rt id
	while IFS=$'\t' read -r rt id || [[ -n $rt ]]; do
		[[ -n $rt && $rt != \#* ]] || continue
		if [[ ! $id =~ ^sha256:[0-9a-f]{64}$ ]]; then
			warn "ignoring malformed line in ${tags_file##*/}: $rt $id"
			continue
		fi
		before_id[$id]=1
		[[ $rt == *'<none>'* ]] || before_tag[$rt]=$id
	done <"$tags_file"
	while read -r id || [[ -n $id ]]; do
		[[ -n $id && $id != \#* ]] || continue
		before_id[$id]=1
	done <"$ids_file"
}

# read_after: (re)load the current image/tag/container state from the daemon.
read_after() {
	local out rt id status name
	local -a cids=()
	tag_id=() tag_count=() id_present=() used_running=() used_stopped=()

	out=$(docker images --no-trunc --format '{{.Repository}}:{{.Tag}}\t{{.ID}}') || die "docker images failed"
	while IFS=$'\t' read -r rt id; do
		[[ -n $rt && $rt != *'<none>'* ]] || continue
		tag_id[$rt]=$id
		tag_count[$id]=$((${tag_count[$id]:-0} + 1))
	done <<<"$out"

	out=$(docker images -a -q --no-trunc) || die "docker images -a failed"
	while read -r id; do
		[[ -z $id ]] || id_present[$id]=1
	done <<<"$out"

	# Containers can vanish between ps and inspect (other --rm runs finishing),
	# so inspect errors for individual containers are ignored.
	out=$(docker ps -a -q --no-trunc) || die "docker ps failed"
	[[ -z $out ]] || mapfile -t cids <<<"$out"
	((${#cids[@]})) || return 0
	out=$(docker inspect --format '{{.Image}}{{"\t"}}{{.State.Status}}{{"\t"}}{{.Name}}' "${cids[@]}" 2>/dev/null) || true
	while IFS=$'\t' read -r id status name; do
		[[ -n $id ]] || continue
		name="${name#/} ($status)"
		case $status in
		running | paused | restarting) used_running[$id]+="${used_running[$id]:+, }$name" ;;
		*) used_stopped[$id]+="${used_stopped[$id]:+, }$name" ;;
		esac
	done <<<"$out"
}

# Model updates, applied after every successful (or, under --dry-run, planned)
# command so later steps see the resulting state.
model_tag() { # ID REF: REF now points at ID
	local id=$1 rt=$2 old=${tag_id[$2]:-}
	[[ -z $old ]] || tag_count[$old]=$((${tag_count[$old]:-1} - 1))
	tag_id[$rt]=$id
	tag_count[$id]=$((${tag_count[$id]:-0} + 1))
	id_present[$id]=1
}
model_untag() { # REF: REF is gone; its image too if it was the last name and unused
	local rt=$1 id=${tag_id[$1]} n
	unset 'tag_id[$rt]'
	n=$((${tag_count[$id]:-1} - 1))
	tag_count[$id]=$n
	if ((n <= 0)) && [[ -z ${before_id[$id]:-} && -z $(users_of "$id") ]]; then
		unset 'id_present[$id]'
	fi
}
model_remove() { # ID: image and all its names are gone
	local id=$1 rt
	unset 'id_present[$id]'
	tag_count[$id]=0
	for rt in "${!tag_id[@]}"; do
		[[ ${tag_id[$rt]} != "$id" ]] || unset 'tag_id[$rt]'
	done
}

# ---- b. re-tag pre-existing names that were taken or vanished ----------------
step_retag() {
	local rt id cur from n=0
	local -a keys
	mapfile -t keys < <(keys_sorted before_tag)
	for rt in "${keys[@]}"; do
		id=${before_tag[$rt]}
		cur=${tag_id[$rt]:-}
		[[ $cur != "$id" ]] || continue
		n=$((n + 1))
		if [[ -z ${id_present[$id]:-} ]] && ! docker image inspect "$id" >/dev/null 2>&1; then
			warn "!!! pre-existing image $(short "$id") for $rt NO LONGER EXISTS: cannot restore $rt (re-pull it by digest or re-snapshot)"
			continue
		fi
		if run docker tag "$id" "$rt"; then
			from="<missing>"
			[[ -z $cur ]] || from=$(short "$cur")
			done_retag+=("$rt: $from -> $(short "$id")")
			model_tag "$id" "$rt"
		else
			failed+=("docker tag $(short "$id") $rt: $run_err")
		fi
	done
	((n)) || echo "  (nothing to do)"
}

# ---- c. untag names that did not exist before -------------------------------
step_untag() {
	local rt id users n=0
	local -a keys
	mapfile -t keys < <(keys_sorted tag_id)
	for rt in "${keys[@]}"; do
		[[ -z ${before_tag[$rt]:-} ]] || continue
		n=$((n + 1))
		id=${tag_id[$rt]}
		if ((${tag_count[$id]:-0} <= 1)); then
			# Last name on this image: `docker rmi` would delete the image as well.
			if [[ -n ${before_id[$id]:-} ]]; then
				skip "tag $rt: only name of pre-existing image $(short "$id"), kept so that image is not deleted"
				continue
			fi
			users=$(users_of "$id")
			if [[ -n $users ]]; then
				skip "tag $rt ($(short "$id")): image in use by $users"
				continue
			fi
		fi
		if run docker rmi --no-prune "$rt"; then
			model_untag "$rt"
			if [[ -n ${id_present[$id]:-} ]]; then
				done_untag+=("$rt ($(short "$id"))")
			else
				done_untag+=("$rt ($(short "$id"), last name: image deleted)")
			fi
		else
			failed+=("docker rmi $rt: $run_err")
		fi
	done
	((n)) || echo "  (nothing to do)"
}

# ---- d. delete image ids that did not exist before ---------------------------
step_remove() {
	local id users label pass nskip=0
	local -a keys todo=() retry=()
	local -A err=()
	mapfile -t keys < <(keys_sorted id_present)
	for id in "${keys[@]}"; do
		[[ -z ${before_id[$id]:-} ]] || continue
		users=$(users_of "$id")
		if [[ -n $users ]]; then
			skip "image $(short "$id")${id_names[$id]:+ (${id_names[$id]})}: in use by $users"
			nskip=$((nskip + 1))
			continue
		fi
		todo+=("$id")
	done
	if ((${#todo[@]} == 0)); then
		((nskip)) || echo "  (nothing to do)"
		return 0
	fi
	# With --no-prune a parent image can only go after its children: retry
	# failures for as long as each pass makes progress.
	for ((pass = 1; ${#todo[@]} > 0; pass++)); do
		retry=()
		for id in "${todo[@]}"; do
			label="$(short "$id")${id_names[$id]:+  ${id_names[$id]}}"
			if ((!dry_run)) && ! docker image inspect "$id" >/dev/null 2>&1; then
				done_remove+=("$label (already gone)")
				continue
			fi
			if run docker rmi -f --no-prune "$id"; then
				done_remove+=("$label")
				model_remove "$id"
			else
				retry+=("$id")
				err[$id]=$run_err
			fi
		done
		((${#retry[@]} < ${#todo[@]})) || break
		todo=("${retry[@]}")
	done
	for id in "${retry[@]}"; do
		failed+=("docker rmi -f $(short "$id"): ${err[$id]}")
	done
}

# section TITLE ITEM...: print a summary section with its item count.
section() {
	local title=$1
	shift
	printf '%-44s %d\n' "$title" "$#"
	(($# == 0)) || printf '    %s\n' "$@"
}

# ---- main -------------------------------------------------------------------
read_before
read_after
while IFS= read -r rt; do
	[[ -z $rt ]] || id_names[${tag_id[$rt]}]+="${id_names[${tag_id[$rt]}]:+ }$rt"
done < <(keys_sorted tag_id)

mode="cleanup"
((!dry_run)) || mode="cleanup DRY RUN (no docker command is executed)"
printf '== %s ==\n' "$mode"
printf 'snapshot: %d tags, %d ids (taken %s)\n' "${#before_tag[@]}" "${#before_id[@]}" \
	"$(date -r "$tags_file" '+%Y-%m-%d %H:%M:%S %Z')"
printf 'now:      %d tags, %d ids, %d image(s) used by containers\n' "${#tag_id[@]}" "${#id_present[@]}" \
	"$(printf '%s\n' "${!used_running[@]}" "${!used_stopped[@]}" | sort -u | grep -c . || true)"

echo "-- b. re-tag pre-existing tags that were taken or vanished"
step_retag
((dry_run)) || read_after
echo "-- c. untag tags that are new since the snapshot"
step_untag
((dry_run)) || read_after
echo "-- d. delete image ids that are new since the snapshot"
step_remove
((dry_run)) || read_after

# ---- e. summary ---------------------------------------------------------------
declare -a bad=() left=() gone=()
while IFS= read -r rt; do
	[[ -n $rt ]] || continue
	id=${before_tag[$rt]}
	cur=${tag_id[$rt]:-}
	[[ $cur != "$id" ]] || continue
	if [[ -z ${id_present[$id]:-} ]]; then
		bad+=("$rt: snapshot image $(short "$id") no longer exists")
	elif [[ -z $cur ]]; then
		bad+=("$rt: tag missing (should be $(short "$id"))")
	else
		bad+=("$rt: points at $(short "$cur"), should be $(short "$id")")
	fi
done < <(keys_sorted before_tag)
while IFS= read -r rt; do
	[[ -n $rt && -z ${before_tag[$rt]:-} ]] || continue
	users=$(users_of "${tag_id[$rt]}")
	left+=("tag   $rt -> $(short "${tag_id[$rt]}")${users:+  [in use by $users]}")
done < <(keys_sorted tag_id)
while IFS= read -r id; do
	[[ -n $id && -z ${before_id[$id]:-} ]] || continue
	users=$(users_of "$id")
	left+=("image $(short "$id")${id_names[$id]:+  ${id_names[$id]}}${users:+  [in use by $users]}")
done < <(keys_sorted id_present)
while IFS= read -r id; do
	[[ -n $id && -z ${id_present[$id]:-} ]] || continue
	gone+=("$(short "$id")")
done < <(keys_sorted before_id)

echo
if ((dry_run)); then
	echo "== summary (DRY RUN: nothing was changed; this is the predicted end state) =="
else
	echo "== summary =="
fi
section "re-tagged (pre-existing tag restored):" "${done_retag[@]}"
section "untagged (new tag removed):" "${done_untag[@]}"
section "removed (new image deleted):" "${done_remove[@]}"
section "skipped:" "${skipped[@]}"
section "failed:" "${failed[@]}"
section "left over (new since snapshot, still here):" "${left[@]}"
section "pre-existing untagged ids gone (info only):" "${gone[@]}"
section "PRE-EXISTING TAGS MISSING OR WRONG:" "${bad[@]}"
if ((${#bad[@]})); then
	echo "result: FAILED - ${#bad[@]} pre-existing repo:tag(s) not restored"
	exit 1
fi
echo "result: OK - all ${#before_tag[@]} pre-existing repo:tags point at their snapshot ids"