#!/usr/bin/env bash
#
# eks-kube-upgrade-rollover.sh
#
# Codifies Steps 2-3 + Verification of the EKS Kubernetes Upgrade Runbook:
#   docs/platform/infrastructure/eks-kube-upgrade-runbook.md
#
# Step 1 (bumping spec.kubernetes.version in the platform-versions repo's
# VersionConfig) is NOT scriptable here — it's a manual PR in a separate,
# non-customer-accessible repo. Run this script only after that PR has
# merged and the new version is live.
#
# Usage:
#   ./eks-kube-upgrade-rollover.sh --target-version 1.32 [options]
#
# Options:
#   --target-version VERSION   Required. Kube minor version nodes should be on, e.g. 1.32
#                               (matched as a prefix against KUBELET_VERSION, e.g. v1.32.*)
#   --karpenter-namespace NS    Namespace of the karpenter deployment (default: karpenter)
#   --drain-timeout DURATION    Passed to `kubectl drain --timeout` (default: 120s)
#   --apply                     Actually run mutating commands. Without this flag, the
#                               script only prints (dry-run) the commands it would run.
#   -h, --help                  Show this help and exit
#
# Safety:
#   - Defaults to dry-run. Mutating phases (karpenter restart, cordon, drain,
#     PDB-hang remediation) only execute with --apply, and --apply itself is
#     gated behind an interactive y/N confirmation naming the current
#     kubectl context/cluster before anything runs.
#   - On a drain that stalls (PodDisruptionBudget violation), the script
#     identifies the pod's owning Deployment/StatefulSet and prompts before
#     doing anything — it never auto-restarts or auto-deletes.

set -euo pipefail

TARGET_VERSION=""
KARPENTER_NAMESPACE="karpenter"
DRAIN_TIMEOUT="120s"
APPLY=false

usage() {
    sed -n '2,32p' "$0" | sed 's/^# \{0,1\}//'
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --target-version)
            TARGET_VERSION="$2"; shift 2 ;;
        --karpenter-namespace)
            KARPENTER_NAMESPACE="$2"; shift 2 ;;
        --drain-timeout)
            DRAIN_TIMEOUT="$2"; shift 2 ;;
        --apply)
            APPLY=true; shift ;;
        -h|--help)
            usage; exit 0 ;;
        *)
            echo "Unknown argument: $1" >&2; usage; exit 1 ;;
    esac
done

if [[ -z "$TARGET_VERSION" ]]; then
    echo "ERROR: --target-version is required (e.g. --target-version 1.32)" >&2
    exit 1
fi

if ! command -v kubectl >/dev/null 2>&1; then
    echo "ERROR: kubectl not found on PATH" >&2
    exit 1
fi

if ! command -v jq >/dev/null 2>&1; then
    echo "ERROR: jq not found on PATH" >&2
    exit 1
fi

log()  { printf '[%s] %s\n' "$(date '+%H:%M:%S')" "$*"; }
plan() { printf '[DRY-RUN] %s\n' "$*"; }

run() {
    if [[ "$APPLY" == true ]]; then
        log "+ $*"
        "$@"
    else
        plan "$*"
    fi
}

CURRENT_CONTEXT="$(kubectl config current-context)"
log "kubectl context: ${CURRENT_CONTEXT}"
log "Target kube minor version: ${TARGET_VERSION}"

SERVER_VERSION="$(kubectl version -o json | jq -r '.serverVersion | "\(.major).\(.minor)"' | tr -d '+')"
if [[ "${SERVER_VERSION}" != "${TARGET_VERSION}" ]]; then
    echo "ERROR: API server is on ${SERVER_VERSION}, not ${TARGET_VERSION}." >&2
    echo "This means the VersionConfig bump (Step 1) for ${TARGET_VERSION} hasn't merged/rolled out yet — target version is invalid." >&2
    exit 1
fi
log "API server confirmed on ${SERVER_VERSION}"

if [[ "$APPLY" == true ]]; then
    echo
    echo "About to mutate cluster reachable via context '${CURRENT_CONTEXT}':"
    echo "  - restart deploy/karpenter -n ${KARPENTER_NAMESPACE}"
    echo "  - cordon and drain every node not on kube version ${TARGET_VERSION}.*"
    read -r -p "Proceed? [y/N] " confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        log "Aborted by user."
        exit 1
    fi
fi

# --- Step 2: Manually restart Karpenter on Fargate ---
log "Checking Fargate node versions"
FARGATE_OLD=$(
    kubectl get nodes -l eks.amazonaws.com/compute-type=fargate -o json |
        jq -r --arg v "v${TARGET_VERSION}." \
            '[.items[] | select(.status.nodeInfo.kubeletVersion | startswith($v) | not)] | length'
)

if [[ "$FARGATE_OLD" -eq 0 ]]; then
    log "All Fargate nodes already on ${TARGET_VERSION}.* — skipping karpenter restart"
else
    log "${FARGATE_OLD} Fargate node(s) not on ${TARGET_VERSION}.* — restarting karpenter deployment so it schedules onto a new-version node"
    run kubectl rollout restart "deploy/karpenter" --namespace "${KARPENTER_NAMESPACE}"
    if [[ "$APPLY" == true ]]; then
        log "Waiting for karpenter rollout to finish"
        kubectl rollout status "deploy/karpenter" --namespace "${KARPENTER_NAMESPACE}" --timeout=120s
    else
        plan "kubectl rollout status deploy/karpenter --namespace ${KARPENTER_NAMESPACE} --timeout=120s"
    fi
fi

# --- Step 3: Identify and roll over old-version Karpenter-managed nodes ---
log "Identifying nodes not on kube version ${TARGET_VERSION}.*"

mapfile -t OLD_NODES < <(
    kubectl get nodes -o json |
        jq -r --arg v "v${TARGET_VERSION}." \
            '.items[] | select(.status.nodeInfo.kubeletVersion | startswith($v) | not) | .metadata.name'
)

if [[ ${#OLD_NODES[@]} -eq 0 ]]; then
    log "No old-version nodes found. Nothing to cordon/drain."
else
    log "Old-version nodes (${#OLD_NODES[@]}): ${OLD_NODES[*]}"

    # Cordon ALL old-version nodes before draining any (see runbook caution:
    # draining node A while old-version node B is still uncordoned can let a
    # pod evicted from A land right back on B).
    log "Cordoning all old-version nodes first"
    for node in "${OLD_NODES[@]}"; do
        run kubectl cordon "${node}"
    done

    log "Draining old-version nodes"
    for node in "${OLD_NODES[@]}"; do
        if [[ "$APPLY" != true ]]; then
            plan "kubectl drain ${node} --ignore-daemonsets --delete-emptydir-data --timeout=${DRAIN_TIMEOUT}"
            continue
        fi

        log "+ kubectl drain ${node} --ignore-daemonsets --delete-emptydir-data --timeout=${DRAIN_TIMEOUT}"
        if kubectl drain "${node}" --ignore-daemonsets --delete-emptydir-data --timeout="${DRAIN_TIMEOUT}"; then
            log "Drained ${node}"
            continue
        fi

        log "Drain stalled on ${node} — likely a PodDisruptionBudget violation. Inspecting remaining pods."
        mapfile -t STUCK_PODS < <(
            kubectl get pods --all-namespaces \
                --field-selector "spec.nodeName=${node}" \
                -o json |
                jq -r '.items[] | select(.metadata.ownerReferences != null and .metadata.ownerReferences[0].kind != "DaemonSet") |
                    "\(.metadata.namespace)\t\(.metadata.name)\t\(.metadata.ownerReferences[0].kind)\t\(.metadata.ownerReferences[0].name)"'
        )

        if [[ ${#STUCK_PODS[@]} -eq 0 ]]; then
            log "No pods with an owning workload found on ${node}; leaving it cordoned for manual follow-up."
            continue
        fi

        for entry in "${STUCK_PODS[@]}"; do
            IFS=$'\t' read -r ns pod owner_kind owner_name <<<"$entry"

            if [[ "$owner_kind" == "ReplicaSet" ]]; then
                rs_owner_kind="$(kubectl get replicaset "${owner_name}" -n "${ns}" \
                    -o jsonpath='{.metadata.ownerReferences[0].kind}' 2>/dev/null || true)"
                rs_owner_name="$(kubectl get replicaset "${owner_name}" -n "${ns}" \
                    -o jsonpath='{.metadata.ownerReferences[0].name}' 2>/dev/null || true)"
                if [[ -n "$rs_owner_kind" && -n "$rs_owner_name" ]]; then
                    owner_kind="$rs_owner_kind"
                    owner_name="$rs_owner_name"
                fi
            fi

            echo
            echo "Stuck pod: ${ns}/${pod}  (owner: ${owner_kind}/${owner_name})"

            if [[ "$owner_kind" != "Deployment" && "$owner_kind" != "StatefulSet" ]]; then
                echo "Owner kind '${owner_kind}' is not a Deployment/StatefulSet — runbook remediation doesn't apply. Handle manually."
                continue
            fi

            ready_count="$(kubectl get "${owner_kind,,}" "${owner_name}" -n "${ns}" \
                -o jsonpath='{.status.readyReplicas}' 2>/dev/null || true)"
            ready_count="${ready_count:-0}"

            if [[ "$ready_count" == "0" ]]; then
                echo "This app appears to have no ready pods anywhere — 'rollout restart' would hang forever (new pod never becomes ready)."
                read -r -p "Delete pod ${ns}/${pod} directly instead? [y/N] " del_confirm
                if [[ "$del_confirm" =~ ^[Yy]$ ]]; then
                    kubectl delete pod "${pod}" -n "${ns}"
                else
                    log "Skipped ${ns}/${pod}; leaving ${node} cordoned for manual follow-up."
                fi
            else
                read -r -p "Restart ${owner_kind}/${owner_name} in ${ns} to force the pod off ${node}? [y/N] " restart_confirm
                if [[ "$restart_confirm" =~ ^[Yy]$ ]]; then
                    kubectl rollout restart "${owner_kind,,}/${owner_name}" -n "${ns}"
                else
                    log "Skipped ${owner_kind}/${owner_name} in ${ns}; leaving ${node} cordoned for manual follow-up."
                fi
            fi
        done
    done
fi

# --- Verification ---
log "Verification: current node versions"
run_verification() {
    kubectl get nodes -o wide
}

if [[ "$APPLY" == true ]]; then
    log "Waiting 120s for node status to settle before verifying"
    sleep 120
    run_verification
    REMAINING=$(kubectl get nodes -o json | jq -r --arg v "v${TARGET_VERSION}." \
        '[.items[] | select(.status.nodeInfo.kubeletVersion | startswith($v) | not)] | length')
    if [[ "$REMAINING" -gt 0 ]]; then
        log "WARNING: ${REMAINING} node(s) still not on ${TARGET_VERSION}.* — repeat this script or wait for TTL rollover before starting the next VersionConfig bump."
    else
        log "All nodes on ${TARGET_VERSION}.*. Safe to proceed to the next VersionConfig minor-version bump."
    fi
else
    plan "kubectl get nodes -o wide"
    plan "(re-check for remaining old-version nodes)"
fi