#!/usr/bin/env bash
set -Eeuo pipefail
IFS=$'\n\t'

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
# Like e2e-k8s-dispatch, default to the invoking checkout's plugin tests/e2e
# directory; a hardcoded worktree path here rots as worktrees come and go.
E2E_DIR="${IPZ_E2E_DIR:-$PWD}"
BUILD_SCRIPT="${IPZ_E2E_RUNNER_BUILD_SCRIPT:-$SCRIPT_DIR/../image/build.sh}"
DEFAULT_NODES=(
    "debian1=100.106.253.50"
    "debian2=100.79.69.43"
    "debian3=100.101.104.41"
)

CHECK_ONLY=0
DRY_RUN=0
PODMAN_BIN=""
IMAGE=""
PLAYWRIGHT_VERSION=""
LOCAL_STATE="unknown"
LOCAL_ARCHIVE=""
NODE_ERRORS=0

# These arrays are populated after parsing IPZ_E2E_K8S_NODES.  Keeping the
# label and address separate makes the output useful without putting host
# aliases into the SSH destination accidentally.
declare -a NODE_LABELS=()
declare -a NODE_ADDRESSES=()
declare -a NODE_STATES=()

default_nodes_value() {
    local node
    local -a values=()
    for node in "${DEFAULT_NODES[@]}"; do
        values+=("$node")
    done
    local IFS=' '
    printf '%s' "${values[*]}"
}

usage() {
    printf 'Usage: %s [--check] [--dry-run]\n' "${BASH_SOURCE[0]}" >&2
    printf '\n' >&2
    printf 'Synchronize the pinned IPZ Playwright runner image to every k3s node.\n' >&2
    printf '  --check       inspect local and node images without building or importing\n' >&2
    printf '  --dry-run     inspect state and print build/import actions without acting\n' >&2
    printf '  --help        show this help\n' >&2
    printf '\n' >&2
    printf 'Environment overrides:\n' >&2
    printf '  IPZ_E2E_DIR             tests/e2e source (default: %s)\n' "$DEFAULT_E2E_DIR" >&2
    printf '  IPZ_E2E_K8S_NODES       space/comma-separated label=address entries (default: %s)\n' "$(default_nodes_value)" >&2
    printf '  IPZ_PODMAN_BIN         podman-compatible executable\n' >&2
    printf '  IPZ_E2E_SSH_KEY        buildbox SSH key (default: %s)\n' "${HOME:-}/.ssh/id_ed25519_buildbox" >&2
    printf '  IPZ_E2E_SSH_PORT       buildbox SSH port (default: 2222)\n' >&2
    printf '  IPZ_E2E_SSH_USER       buildbox SSH user (default: user)\n' >&2
    printf '  IPZ_E2E_SCP_BIN        scp-compatible executable for archive delivery\n' >&2
    printf '  IPZ_E2E_RUNNER_BUILD_SCRIPT  image/build.sh override\n' >&2
}

fail() {
    printf 'e2e-runner-image-sync: %s\n' "$1" >&2
    exit 1
}

on_exit() {
    local status=$?
    if [[ -n "$LOCAL_ARCHIVE" ]]; then
        rm -f -- "$LOCAL_ARCHIVE" || true
    fi
    if ((status != 0)); then
        printf 'e2e-runner-image-sync: failed (exit %d)\n' "$status" >&2
    fi
}
trap on_exit EXIT

while (($# > 0)); do
    case "$1" in
        --check)
            CHECK_ONLY=1
            ;;
        --dry-run)
            DRY_RUN=1
            ;;
        --help|-h)
            usage
            exit 0
            ;;
        *)
            usage
            fail "unknown option: $1"
            ;;
    esac
    shift
done

[[ -d "$E2E_DIR" ]] || fail "E2E directory does not exist: $E2E_DIR (set IPZ_E2E_DIR or run from the plugin tests/e2e directory)"
[[ -r "$E2E_DIR/package.json" ]] || fail "package.json is not readable: $E2E_DIR/package.json"
[[ -r "$E2E_DIR/package-lock.json" ]] || fail "package-lock.json is not readable: $E2E_DIR/package-lock.json"
[[ -x "$BUILD_SCRIPT" ]] || fail "image build script is not executable: $BUILD_SCRIPT"
command -v python3 >/dev/null 2>&1 || fail "python3 is required to validate the lockfile"

PLAYWRIGHT_VERSION="$(python3 - "$E2E_DIR/package.json" "$E2E_DIR/package-lock.json" <<'PY'
import json
import re
import sys

package_path, lock_path = sys.argv[1:]
with open(package_path, encoding="utf-8") as stream:
    package = json.load(stream)
with open(lock_path, encoding="utf-8") as stream:
    lock = json.load(stream)

version = package.get("devDependencies", {}).get("@playwright/test")
if not isinstance(version, str) or not re.fullmatch(r"\d+\.\d+\.\d+", version):
    raise SystemExit("@playwright/test must be an exact semver in package.json")

packages = lock.get("packages", {})
for package_name in ("@playwright/test", "playwright", "playwright-core"):
    lock_entry = packages.get(f"node_modules/{package_name}")
    if not isinstance(lock_entry, dict) or lock_entry.get("version") != version:
        actual = lock_entry.get("version") if isinstance(lock_entry, dict) else None
        raise SystemExit(
            f"{package_name} is {actual!r} in package-lock.json, expected {version!r}"
        )

print(version)
PY
)"
IMAGE="localhost/ipz-e2e-runner:${PLAYWRIGHT_VERSION}"

if [[ -n "${IPZ_PODMAN_BIN:-}" ]]; then
    PODMAN_BIN="$IPZ_PODMAN_BIN"
elif command -v deck-podman >/dev/null 2>&1; then
    PODMAN_BIN="$(command -v deck-podman)"
elif command -v podman >/dev/null 2>&1; then
    PODMAN_BIN="$(command -v podman)"
else
    fail "podman is required"
fi
[[ -x "$PODMAN_BIN" ]] || fail "Podman runner is not executable: $PODMAN_BIN"

SSH_BIN="${IPZ_E2E_SSH_BIN:-ssh}"
SCP_BIN="${IPZ_E2E_SCP_BIN:-scp}"
SSH_KEY="${IPZ_E2E_SSH_KEY:-${HOME:-}/.ssh/id_ed25519_buildbox}"
SSH_PORT="${IPZ_E2E_SSH_PORT:-2222}"
SSH_USER="${IPZ_E2E_SSH_USER:-user}"
[[ "$SSH_PORT" =~ ^[0-9]+$ ]] || fail "IPZ_E2E_SSH_PORT must be numeric"
command -v "$SSH_BIN" >/dev/null 2>&1 || fail "SSH executable is not available: $SSH_BIN"
[[ -r "$SSH_KEY" ]] || fail "SSH key is not readable: $SSH_KEY"

# The default is the audited three-node buildbox set.  Custom values are
# intentionally opt-in for local tests and isolated clusters.
NODE_SPEC="${IPZ_E2E_K8S_NODES:-$(default_nodes_value)}"
NODE_SPEC="${NODE_SPEC//,/ }"
IFS=' ' read -r -a NODE_ENTRIES <<< "$NODE_SPEC"
((${#NODE_ENTRIES[@]} > 0)) || fail "IPZ_E2E_K8S_NODES must contain at least one node"

for entry in "${NODE_ENTRIES[@]}"; do
    [[ "$entry" == *=* ]] || fail "node entry must be label=address: $entry"
    label="${entry%%=*}"
    address="${entry#*=}"
    [[ -n "$label" && -n "$address" ]] || fail "node entry has an empty label or address: $entry"
    [[ "$label" != *[[:space:]]* && "$address" != *[[:space:]]* ]] || fail "node entry contains whitespace: $entry"
    NODE_LABELS+=("$label")
    NODE_ADDRESSES+=("$address")
    NODE_STATES+=("unknown")
done

printf 'Playwright version: %s\n' "$PLAYWRIGHT_VERSION"
printf 'Runner image: %s\n' "$IMAGE"
printf 'Podman runner: %s\n' "$PODMAN_BIN"
if ((CHECK_ONLY)); then
    printf 'Mode: check (read-only)\n'
elif ((DRY_RUN)); then
    printf 'Mode: dry-run (no build or import)\n'
else
    printf 'Mode: synchronize\n'
fi

local_image_present() {
    local status
    if "$PODMAN_BIN" image exists "$IMAGE" >/dev/null 2>&1; then
        LOCAL_STATE="present"
        return 0
    else
        status=$?
    fi
    if ((status != 1)); then
        fail "podman image exists failed for $IMAGE (exit $status)"
    fi
    LOCAL_STATE="absent"
    return 1
}

if local_image_present; then
    printf 'Local image: present\n'
else
    printf 'Local image: absent\n'
fi

ssh_node_image_listing() {
    local address="$1"
    "$SSH_BIN" \
        -p "$SSH_PORT" \
        -i "$SSH_KEY" \
        -o BatchMode=yes \
        -o ConnectTimeout="${IPZ_E2E_SSH_CONNECT_TIMEOUT:-15}" \
        "$SSH_USER@$address" bash -s <<'REMOTE'
set -Eeuo pipefail
sudo /usr/local/bin/k3s ctr -n k8s.io images ls
REMOTE
}

listing_contains_image() {
    local listing="$1"
    LC_ALL=C awk -v expected="$IMAGE" '$1 == expected { found = 1 } END { exit(found ? 0 : 1) }' <<< "$listing"
}

query_node() {
    local index="$1"
    local label="${NODE_LABELS[$index]}"
    local address="${NODE_ADDRESSES[$index]}"
    local listing

    if ! listing="$(ssh_node_image_listing "$address")"; then
        NODE_STATES[index]="error"
        NODE_ERRORS=$((NODE_ERRORS + 1))
        printf 'Node %s (%s): error (unable to list k3s images)\n' "$label" "$address" >&2
        return 0
    fi

    if listing_contains_image "$listing"; then
        NODE_STATES[index]="present"
        printf 'Node %s (%s): present\n' "$label" "$address"
    else
        NODE_STATES[index]="absent"
        printf 'Node %s (%s): absent\n' "$label" "$address"
    fi
}

for ((index = 0; index < ${#NODE_LABELS[@]}; index++)); do
    query_node "$index"
done

((NODE_ERRORS == 0)) || fail "unable to inspect $NODE_ERRORS k3s node(s)"

missing_nodes=()
for ((index = 0; index < ${#NODE_LABELS[@]}; index++)); do
    if [[ "${NODE_STATES[$index]}" == "absent" ]]; then
        missing_nodes+=("$index")
    fi
done

if ((CHECK_ONLY)); then
    if [[ "$LOCAL_STATE" == "present" && ${#missing_nodes[@]} -eq 0 ]]; then
        printf 'Check result: synchronized\n'
        exit 0
    fi
    printf 'Check result: synchronization required (local=%s, missing_nodes=%d)\n' \
        "$LOCAL_STATE" "${#missing_nodes[@]}"
    exit 1
fi

if ((DRY_RUN)); then
    if [[ "$LOCAL_STATE" == "absent" ]]; then
        printf 'Plan: build %s with %s\n' "$IMAGE" "$BUILD_SCRIPT"
    else
        printf 'Plan: local image already exists; skip build\n'
    fi
    if ((${#missing_nodes[@]} == 0)); then
        printf 'Plan: all nodes already contain %s; no imports\n' "$IMAGE"
    else
        for index in "${missing_nodes[@]}"; do
            printf 'Plan: import %s to node %s (%s)\n' "$IMAGE" "${NODE_LABELS[$index]}" "${NODE_ADDRESSES[$index]}"
        done
    fi
    printf 'Dry-run complete: no image build or import performed.\n'
    exit 0
fi

if [[ "$LOCAL_STATE" == "absent" ]]; then
    printf 'Building %s with %s\n' "$IMAGE" "$BUILD_SCRIPT"
    IPZ_PODMAN_BIN="$PODMAN_BIN" "$BUILD_SCRIPT"
    if ! local_image_present; then
        fail "image build completed but $IMAGE is still absent locally"
    fi
    printf 'Local image: present after build\n'
fi

ensure_local_archive() {
    if [[ -n "$LOCAL_ARCHIVE" ]]; then
        return 0
    fi

    command -v "$SCP_BIN" >/dev/null 2>&1 || fail "SCP executable is not available: $SCP_BIN"
    LOCAL_ARCHIVE="$(mktemp "${TMPDIR:-/tmp}/ipz-e2e-runner.XXXXXX.tar")" || fail "unable to create a local image archive"
    printf 'Saving %s to local archive\n' "$IMAGE"
    if ! "$PODMAN_BIN" save "$IMAGE" >"$LOCAL_ARCHIVE"; then
        fail "unable to save $IMAGE to $LOCAL_ARCHIVE"
    fi
    [[ -s "$LOCAL_ARCHIVE" ]] || fail "podman save produced an empty archive for $IMAGE"
}

cleanup_remote_archive() {
    local address="$1"
    local remote_archive="$2"

    "$SSH_BIN" \
        -p "$SSH_PORT" \
        -i "$SSH_KEY" \
        -o BatchMode=yes \
        -o ConnectTimeout="${IPZ_E2E_SSH_CONNECT_TIMEOUT:-15}" \
        "$SSH_USER@$address" bash -s -- "$remote_archive" <<'REMOTE'
set -Eeuo pipefail
rm -f -- "$1"
REMOTE
}

ship_to_node() {
    local index="$1"
    local label="${NODE_LABELS[$index]}"
    local address="${NODE_ADDRESSES[$index]}"
    local remote_archive="/tmp/ipz-e2e-runner-${PLAYWRIGHT_VERSION}-$$-${index}.tar"

    ensure_local_archive
    printf 'Copying image archive to node %s (%s)\n' "$label" "$address"
    if ! "$SCP_BIN" \
        -P "$SSH_PORT" \
        -i "$SSH_KEY" \
        -o BatchMode=yes \
        -o ConnectTimeout="${IPZ_E2E_SSH_CONNECT_TIMEOUT:-15}" \
        "$LOCAL_ARCHIVE" \
        "$SSH_USER@$address:$remote_archive"; then
        cleanup_remote_archive "$address" "$remote_archive" || true
        fail "image archive transfer failed on node $label ($address)"
    fi

    printf 'Importing %s to node %s (%s)\n' "$IMAGE" "$label" "$address"
    if ! "$SSH_BIN" \
        -p "$SSH_PORT" \
        -i "$SSH_KEY" \
        -o BatchMode=yes \
        -o ConnectTimeout="${IPZ_E2E_SSH_CONNECT_TIMEOUT:-15}" \
        "$SSH_USER@$address" bash -s -- "$remote_archive" "$IMAGE" <<'REMOTE'
set -Eeuo pipefail
archive_path="$1"
image="$2"
cleanup() {
    rm -f -- "$archive_path"
}
trap cleanup EXIT
sudo /usr/local/bin/k3s ctr -n k8s.io images import "$archive_path"
# Pin in the SAME session as the import: kubelet image GC on these
# disk-pressured nodes has pruned a fresh import before a separate pin
# connection could label it.
sudo /usr/local/bin/k3s ctr -n k8s.io images label "$image" io.cri-containerd.pinned=pinned >/dev/null
REMOTE
    then
        cleanup_remote_archive "$address" "$remote_archive" || true
        fail "image import failed on node $label ($address)"
    fi
}

pin_on_node() {
    # Kubelet image GC prunes any unpinned image no pod is using once node
    # disk crosses its threshold — the buildboxes hover right at it, so an
    # imported runner image vanishes between sync and dispatch. CRI-pinned
    # images are exempt; the label set is idempotent.
    local index="$1"
    local label="${NODE_LABELS[$index]}"
    local address="${NODE_ADDRESSES[$index]}"

    printf 'Pinning %s on node %s (%s)\n' "$IMAGE" "$label" "$address"
    if ! "$SSH_BIN" \
        -p "$SSH_PORT" \
        -i "$SSH_KEY" \
        -o BatchMode=yes \
        -o ConnectTimeout="${IPZ_E2E_SSH_CONNECT_TIMEOUT:-15}" \
        "$SSH_USER@$address" bash -s -- "$IMAGE" <<'REMOTE'
set -Eeuo pipefail
sudo /usr/local/bin/k3s ctr -n k8s.io images label "$1" io.cri-containerd.pinned=pinned >/dev/null
REMOTE
    then
        printf 'Unable to pin %s against kubelet image GC on node %s (%s)\n' "$IMAGE" "$label" "$address" >&2
        return 1
    fi
}

for index in "${missing_nodes[@]}"; do
    ship_to_node "$index"
    query_node "$index"
    [[ "${NODE_STATES[$index]}" == "present" ]] || fail "image import did not make $IMAGE present on node ${NODE_LABELS[$index]}"
done

# Freshly shipped nodes were pinned atomically with their import. Nodes that
# already held the image still need the (idempotent) pin — and the image may
# have been GC-pruned since the presence check, so a failed pin re-ships once.
for index in "${!NODE_LABELS[@]}"; do
    if pin_on_node "$index"; then
        continue
    fi
    printf 'Pin failed on node %s — re-importing before one retry\n' "${NODE_LABELS[$index]}"
    ship_to_node "$index"
    query_node "$index"
    [[ "${NODE_STATES[$index]}" == "present" ]] || fail "image re-import did not make $IMAGE present on node ${NODE_LABELS[$index]}"
    pin_on_node "$index" || fail "unable to pin $IMAGE on node ${NODE_LABELS[$index]} after re-import"
done

printf 'Synchronization complete: %s is present and GC-pinned on all %d node(s).\n' "$IMAGE" "${#NODE_LABELS[@]}"