#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
source "$ROOT/tests/kind/lib.sh"

USER_KUBECONFIG_BEFORE="$(user_kubeconfig_digest)"
SUCCEEDED=false
OWNS_FIXTURE=false
APP_IMAGE=nuthatch-recovery:local

cleanup() {
  local status=$?
  trap - EXIT INT TERM
  if [[ "${NUTHATCH_TESTBED_RETAIN:-0}" = 1 ]] && [[ "$SUCCEEDED" != true ]]; then
    printf 'retained kubeconfig: %s\n' "$KUBECONFIG_PATH"
    printf 'retained context: %s\n' "$CONTEXT"
    exit "$status"
  fi
  if [[ "$OWNS_FIXTURE" = true ]]; then
    destroy_fixture || status=1
  fi
  if [[ "$(user_kubeconfig_digest)" != "$USER_KUBECONFIG_BEFORE" ]]; then
    printf 'user kubeconfig changed unexpectedly\n' >&2
    status=1
  fi
  exit "$status"
}

trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

pg_value() {
  local database="$1"
  local query="$2"
  kube --namespace "$FIXTURE_NAMESPACE" exec deployment/postgres -- \
    env PGPASSWORD=synthetic-postgres psql \
      --username nuthatch \
      --dbname "$database" \
      --tuples-only \
      --no-align \
      --set ON_ERROR_STOP=1 \
      --command "$query" |
    tr -d '[:space:]'
}

wait_for_value() {
  local database="$1"
  local query="$2"
  local expected="$3"
  local description="$4"
  for _ in $(seq 1 120); do
    if [[ "$(pg_value "$database" "$query" 2>/dev/null || true)" = "$expected" ]]; then
      return
    fi
    sleep 0.25
  done
  printf 'timed out waiting for %s\n' "$description" >&2
  return 1
}

render_application() {
  local lane="$1"
  local profile_file="$2"
  local pg_conn
  local redis_url
  local s3_endpoint
  local s3_region
  local s3_bucket
  local s3_access_key
  local s3_secret_key
  pg_conn="$(sed -n 's/^NUTHATCH_PG_CONN=//p' "$profile_file")"
  redis_url="$(sed -n 's/^NUTHATCH_REDIS_URL=//p' "$profile_file")"
  s3_endpoint="$(sed -n 's/^NUTHATCH_S3_ENDPOINT=//p' "$profile_file")"
  s3_region="$(sed -n 's/^NUTHATCH_S3_REGION=//p' "$profile_file")"
  s3_bucket="$(sed -n 's/^NUTHATCH_S3_BUCKET=//p' "$profile_file")"
  s3_access_key="$(sed -n 's/^NUTHATCH_S3_ACCESS_KEY=//p' "$profile_file")"
  s3_secret_key="$(sed -n 's/^NUTHATCH_S3_SECRET_KEY=//p' "$profile_file")"
  sed \
    -e "s|@@LANE@@|$lane|g" \
    -e "s|@@APP_IMAGE@@|$APP_IMAGE|g" \
    -e "s|@@PG_CONN@@|$pg_conn|g" \
    -e "s|@@REDIS_URL@@|$redis_url|g" \
    -e "s|@@OIDC_ISSUER@@|https://issuer.invalid|g" \
    -e "s|@@OIDC_CLIENT_ID@@|recovery|g" \
    -e "s|@@OIDC_ADDITIONAL_AUDIENCES@@||g" \
    -e "s|@@OIDC_CLIENT_SECRET@@|synthetic-recovery-oidc-client-secret|g" \
    -e "s|@@S3_ENDPOINT@@|$s3_endpoint|g" \
    -e "s|@@S3_REGION@@|$s3_region|g" \
    -e "s|@@S3_BUCKET@@|$s3_bucket|g" \
    -e "s|@@S3_ACCESS_KEY@@|$s3_access_key|g" \
    -e "s|@@S3_SECRET_KEY@@|$s3_secret_key|g" \
    "$ROOT/deploy/kubernetes/nuthatch.yaml.in" >"$RUNTIME/recovery-app-$lane.yaml"
  sed -i \
    's/NUTHATCH_OBJECT_REQUEST_TIMEOUT_SEC: "300"/NUTHATCH_OBJECT_REQUEST_TIMEOUT_SEC: "2"/' \
    "$RUNTIME/recovery-app-$lane.yaml"
  sed -i \
    '/NUTHATCH_OBJECT_REQUEST_TIMEOUT_SEC/a\  NUTHATCH_RECOVERY_GRACE_SEC: "1"' \
    "$RUNTIME/recovery-app-$lane.yaml"
  sed -i 's/nameservers: \[10.96.0.53\]/nameservers: [10.96.0.10]/' \
    "$RUNTIME/recovery-app-$lane.yaml"
  sed -i 's/  replicas: 1/  replicas: 2/' "$RUNTIME/recovery-app-$lane.yaml"
}

deploy_recovery() {
  local lane="$1"
  local profile="$2"
  render_application "$lane" "$RUNTIME/$profile.env"
  kube apply --filename "$RUNTIME/recovery-app-$lane.yaml" >/dev/null
  kube --namespace "$FIXTURE_NAMESPACE" rollout status \
    "deployment/nuthatch-$lane" --timeout=180s
}

cleanup_recovery() {
  local lane="$1"
  kube --namespace "$FIXTURE_NAMESPACE" delete \
    "deployment/nuthatch-$lane" \
    "service/nuthatch-$lane" \
    "configmap/nuthatch-$lane" \
    "secret/nuthatch-$lane" \
    --ignore-not-found >/dev/null
  kube --namespace "$FIXTURE_NAMESPACE" wait \
    --for=delete pod \
    --selector="app=nuthatch-$lane" \
    --timeout=60s >/dev/null
  rm -f "$RUNTIME/recovery-app-$lane.yaml"
}

start_phase() {
  local lane="$1"
  local profile="$2"
  local phase="$3"
  local job="recovery-$phase-$lane"
  render_lane_job \
    "$job" \
    "$TESTBED/recovery-phase-job.yaml.in" \
    "$APP_IMAGE" \
    "$profile"
  sed -i "s|@@PHASE@@|$phase|g" "$RUNTIME/$job.yaml"
  kube apply --filename "$RUNTIME/$job.yaml" >/dev/null
}

kill_phase() {
  local lane="$1"
  local phase="$2"
  local job="recovery-$phase-$lane"
  kube --namespace "$FIXTURE_NAMESPACE" delete pod \
    --selector="job-name=$job" \
    --grace-period=0 \
    --force >/dev/null
  kube --namespace "$FIXTURE_NAMESPACE" delete "job/$job" \
    --ignore-not-found >/dev/null
}

wait_for_phase_checkpoint() {
  local lane="$1"
  local phase="$2"
  local operation="recovery-$phase"
  local job="recovery-$phase-$lane"
  for _ in $(seq 1 120); do
    if kube --namespace "$FIXTURE_NAMESPACE" logs "job/$job" 2>/dev/null |
        grep -F "recovery checkpoint $operation" >/dev/null; then
      return
    fi
    sleep 0.25
  done
  printf 'timed out waiting for %s checkpoint\n' "$phase" >&2
  return 1
}

assert_terminal_authority() {
  local database="$1"
  local operation="$2"
  local result="$3"
  wait_for_value \
    "$database" \
    "SELECT state FROM mutations WHERE operation_id = '$operation'" \
    "$result" \
    "$operation terminal result"
  test "$(
    pg_value "$database" \
      "SELECT count(*) FROM key_locks WHERE operation_id = '$operation'"
  )" = 0
  test "$(
    pg_value "$database" \
      "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
  )" = 0
}

set_service_available() {
  local service="$1"
  local available="$2"
  local selector="$service"
  if [[ "$available" = false ]]; then
    selector=recovery-unavailable
  fi
  kube --namespace "$FIXTURE_NAMESPACE" patch service "$service" \
    --type merge \
    --patch "{\"spec\":{\"selector\":{\"app\":\"$selector\"}}}" >/dev/null
}

terminate_database_connections() {
  local database="$1"
  kube --namespace "$FIXTURE_NAMESPACE" exec deployment/postgres -- \
    env PGPASSWORD=synthetic-postgres psql \
      --username nuthatch \
      --dbname postgres \
      --set ON_ERROR_STOP=1 \
      --command \
        "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$database' AND pid <> pg_backend_pid()" \
      >/dev/null
}

assert_object_absent() {
  local target="$1"
  local bucket="$2"
  local key="$3"
  local endpoint
  endpoint="$(target_loopback_endpoint "$target")"
  if kube --namespace "$FIXTURE_NAMESPACE" exec "deployment/$target" \
    --container s3-readiness -- \
    aws --endpoint-url "$endpoint" s3api head-object \
      --bucket "$bucket" \
      --key "$key" >/dev/null 2>&1; then
    return 1
  fi
}

assert_object_present() {
  local target="$1"
  local bucket="$2"
  local key="$3"
  local endpoint
  endpoint="$(target_loopback_endpoint "$target")"
  kube --namespace "$FIXTURE_NAMESPACE" exec "deployment/$target" \
    --container s3-readiness -- \
    aws --endpoint-url "$endpoint" s3api head-object \
      --bucket "$bucket" \
      --key "$key" >/dev/null
}

assert_cache_invalidated() {
  local lane="$1"
  local value
  value="$(
    kube --namespace "$FIXTURE_NAMESPACE" exec deployment/redis -- \
      redis-cli --raw -n "$((10#$lane))" \
      GET nuthatch:list:recovery-user:shared:: |
      tr -d '[:space:]'
  )"
  test -z "$value"
}

run_phase_recovery() {
  local lane="$1"
  local profile="$2"
  local database="$3"
  local target="$4"
  local phase="$5"
  local durable_state="$6"
  local terminal_state="$7"
  local operation="recovery-$phase"
  start_phase "$lane" "$profile" "$phase"
  wait_for_value \
    "$database" \
    "SELECT state || ':' || external_state FROM mutations WHERE operation_id = '$operation'" \
    "$durable_state" \
    "$phase durable checkpoint"
  if [[ "$phase" = during-put ]]; then
    sleep 1
  fi
  kill_phase "$lane" "$phase"
  deploy_recovery "$lane" "$profile"
  assert_terminal_authority "$database" "$operation" "$terminal_state"
  cleanup_recovery "$lane"
  if [[ "$terminal_state" = succeeded ]]; then
    assert_cache_invalidated "$lane"
  fi
}

acquire_fixture_lock
OWNS_FIXTURE=true
setup_fixture durable coordination target
source "$RUNTIME/fixture.env"
"$NUTHATCH_REAL_NERDCTL" --namespace "$NERDCTL_NAMESPACE" build \
  --file "$ROOT/Dockerfile" \
  --target test-runtime \
  --tag "$APP_IMAGE" \
  "$ROOT"
load_local_image "$APP_IMAGE" "$NERDCTL_NAMESPACE"

while IFS=$'\t' read -r lane coordination oidc target; do
  if [[ -n "${NUTHATCH_RECOVERY_TARGET:-}" ]] &&
      [[ "$target" != "$NUTHATCH_RECOVERY_TARGET" ]]; then
    continue
  fi
  profile="$(write_lane_profile "$lane" redis - "$target")"
  database="nuthatch_lane_$lane"
  bucket="nuthatch-lane-$lane"
  printf 'Running recovery lane %s: target=%s\n' "$lane" "$target"

  run_phase_recovery \
    "$lane" "$profile" "$database" "$target" \
    before-dispatch prepared:idle failed
  run_phase_recovery \
    "$lane" "$profile" "$database" "$target" \
    during-put applying:pending failed
  sleep 5
  assert_object_absent "$target" "$bucket" \
    shared/recovery-during-put.bin

  start_phase "$lane" "$profile" live-timeout
  wait_for_phase_checkpoint "$lane" live-timeout
  test "$(
    pg_value "$database" \
      "SELECT count(*) FROM key_locks WHERE operation_id = 'recovery-live-timeout'"
  )" = 1
  test "$(
    pg_value "$database" \
      "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
  )" = "$((20 * 1024 * 1024))"
  kill_phase "$lane" live-timeout
  deploy_recovery "$lane" "$profile"
  assert_terminal_authority "$database" recovery-live-timeout failed
  cleanup_recovery "$lane"
  sleep 5
  assert_object_absent "$target" "$bucket" \
    shared/recovery-live-timeout.bin

  start_phase "$lane" "$profile" after-dispatch
  wait_for_value \
    "$database" \
    "SELECT state || ':' || external_state FROM mutations WHERE operation_id = 'recovery-after-dispatch'" \
    applying:pending \
    "after-dispatch checkpoint"
  kill_phase "$lane" after-dispatch
  pg_value \
    "$database" \
    "UPDATE mutations SET request_deadline_at = now() + interval '60 seconds', worker_lease_until = now() + interval '60 seconds' WHERE operation_id = 'recovery-after-dispatch'; SELECT 1" \
    >/dev/null
  deploy_recovery "$lane" "$profile"
  before_pg_restarts="$(
    kube --namespace "$FIXTURE_NAMESPACE" get pods \
      --selector="app=nuthatch-$lane" \
      --output json |
      jq '[.items[].status.containerStatuses[].restartCount] | add'
  )"
  set_service_available postgres false
  terminate_database_connections "$database"
  sleep 6
  test "$(
    kube --namespace "$FIXTURE_NAMESPACE" get pods \
      --selector="app=nuthatch-$lane" \
      --output json |
      jq '[.items[].status.containerStatuses[].restartCount] | add'
  )" = "$before_pg_restarts"
  set_service_available postgres true
  wait_for_value "$database" "SELECT 1" 1 "PostgreSQL recovery"
  pg_value \
    "$database" \
    "UPDATE mutations SET request_deadline_at = now() - interval '1 second', worker_lease_until = now() - interval '1 second' WHERE operation_id = 'recovery-after-dispatch'; SELECT 1" \
    >/dev/null
  assert_terminal_authority "$database" recovery-after-dispatch succeeded
  cleanup_recovery "$lane"
  assert_cache_invalidated "$lane"
  assert_object_present "$target" "$bucket" \
    shared/recovery-after-dispatch.bin

  start_phase "$lane" "$profile" after-settlement
  wait_for_value \
    "$database" \
    "SELECT state || ':' || external_state FROM mutations WHERE operation_id = 'recovery-after-settlement'" \
    applying:settled \
    "after-settlement checkpoint"
  kill_phase "$lane" after-settlement
  set_service_available "$target" false
  deploy_recovery "$lane" "$profile"
  sleep 3
  kube --namespace "$FIXTURE_NAMESPACE" logs \
    --selector="app=nuthatch-$lane" \
    --prefix=true |
    grep -F '"result":"exhausted"'
  set_service_available "$target" true
  assert_terminal_authority "$database" recovery-after-settlement succeeded
  cleanup_recovery "$lane"
  assert_cache_invalidated "$lane"

  run_phase_recovery \
    "$lane" "$profile" "$database" "$target" \
    delete-settled applying:settled succeeded
  assert_object_absent "$target" "$bucket" \
    shared/recovery-delete-settled.bin

  cleanup_lane_state "$lane" redis - "$target"
  for phase in \
    before-dispatch \
    during-put \
    live-timeout \
    after-dispatch \
    after-settlement \
    delete-settled; do
    cleanup_lane_resources "recovery-$phase-$lane" "$profile"
  done
done < <(expand_fixture_matrix target)

SUCCEEDED=true
printf 'Nuthatch recovery gate passed\n'