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

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"

drill=${1:-}
confirmation=${2:-}
[[ $confirmation == --execute ]] || {
  echo 'usage: failure-drill.sh {api|database|witness} --execute' >&2
  exit 64
}
load_cluster_env

case $drill in
  api)
    started=$(date +%s)
    pod=$(cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl get pods -n easyai -l app.kubernetes.io/name=easyai-api,easyai.io/site=ningbo -o jsonpath='{.items[0].metadata.name}'")
    cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl delete pod '$pod' -n easyai --wait=false >/dev/null"
    for _ in $(seq 1 20); do
      if curl -fsS --max-time 3 --resolve "ai.51easyai.com:443:${CLUSTER_NINGBO_HOST#*@}" \
        https://ai.51easyai.com/api/v1/readyz | grep -q '"ok":true'; then
        elapsed=$(( $(date +%s) - started ))
        break
      fi
      sleep 1
    done
    [[ ${elapsed:-99} -le 10 ]]
    echo "api_failure_drill=PASS failover_seconds=$elapsed"
    ;;
  database)
    old_primary=$(cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl get cluster easyai-postgres -n easyai -o jsonpath='{.status.currentPrimary}'")
    started=$(date +%s)
    cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl delete pod '$old_primary' -n easyai --grace-period=0 --force >/dev/null"
    for _ in $(seq 1 150); do
      new_primary=$(cluster_ssh "$CLUSTER_NINGBO_HOST" \
        "k3s kubectl get cluster easyai-postgres -n easyai -o jsonpath='{.status.currentPrimary}'")
      if [[ -n $new_primary && $new_primary != "$old_primary" ]]; then
        elapsed=$(( $(date +%s) - started ))
        break
      fi
      sleep 2
    done
    [[ ${elapsed:-999} -le 300 ]]
    curl -fsS --max-time 10 https://ai.51easyai.com/api/v1/readyz | grep -q '"ok":true'
    cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$new_primary" <<'REMOTE'
set -euo pipefail
primary=$1
cleanup_probe() {
  k3s kubectl exec -n easyai "$primary" -c postgres -- \
    psql -X -U postgres -d easyai_ai_gateway -c \
    "DELETE FROM system_settings WHERE setting_key='ha_acceptance_probe';" >/dev/null 2>&1 || true
}
trap cleanup_probe EXIT HUP INT TERM
k3s kubectl exec -n easyai "$primary" -c postgres -- \
  psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway -c \
  "INSERT INTO system_settings(setting_key,value) VALUES('ha_acceptance_probe',jsonb_build_object('at',now())) ON CONFLICT(setting_key) DO UPDATE SET value=excluded.value,updated_at=now();" >/dev/null
for _ in $(seq 1 900); do
  if [[ $(k3s kubectl get cluster easyai-postgres -n easyai -o jsonpath='{.status.readyInstances}') == 2 ]]; then
    replication=$(k3s kubectl exec -n easyai "$primary" -c postgres -- \
      psql -X -U postgres -d easyai_ai_gateway -At -c \
      "SELECT count(*) FROM pg_stat_replication WHERE sync_state IN ('sync','quorum');")
    [[ $replication == 1 ]] && break
  fi
  sleep 2
done
[[ ${replication:-0} == 1 ]]
k3s kubectl exec -n easyai "$primary" -c postgres -- \
  psql -X -v ON_ERROR_STOP=1 -U postgres -d easyai_ai_gateway -c \
  "DELETE FROM system_settings WHERE setting_key='ha_acceptance_probe';" >/dev/null
trap - EXIT HUP INT TERM
REMOTE
    echo "database_failure_drill=PASS promotion_seconds=$elapsed old_primary=$old_primary new_primary=$new_primary"
    ;;
  witness)
    cluster_ssh "$CLUSTER_LOS_ANGELES_HOST" 'systemctl stop k3s'
    trap 'cluster_ssh "$CLUSTER_LOS_ANGELES_HOST" "systemctl start k3s" || true' EXIT HUP INT TERM
    cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl get --raw='/readyz?verbose' | grep -q '\\[+\\]etcd ok'"
    curl -fsS --max-time 10 https://ai.51easyai.com/api/v1/readyz | grep -q '"ok":true'
    cluster_ssh "$CLUSTER_LOS_ANGELES_HOST" 'systemctl start k3s'
    cluster_ssh "$CLUSTER_NINGBO_HOST" \
      'k3s kubectl wait --for=condition=Ready node/easyai-los-angeles --timeout=180s'
    trap - EXIT HUP INT TERM
    echo 'witness_failure_drill=PASS quorum_retained=true'
    ;;
  *)
    echo 'drill must be api, database, or witness' >&2
    exit 64
    ;;
esac