#!/usr/bin/env bash
# The only supported q4.service restart path: drain clients, then require the
# gateway and its authoritative OpenQ4 map to become ready again.
set -u

URL="${Q4_HEALTHZ:-http://127.0.0.1:3280/healthz}"
SERVICE="${Q4_SERVICE:-q4}"
DRAIN_POLLS="${1:-${Q4_RESTART_DRAIN_POLLS:-60}}"
DRAIN_INTERVAL="${Q4_RESTART_DRAIN_INTERVAL_SECONDS:-10}"
READY_TIMEOUT="${Q4_RESTART_READY_TIMEOUT_SECONDS:-120}"
READY_INTERVAL="${Q4_RESTART_READY_INTERVAL_SECONDS:-2}"
CURL_TIMEOUT="${Q4_RESTART_CURL_TIMEOUT_SECONDS:-5}"

positive_integer() {
  case "$2" in
    ''|*[!0-9]*|0) echo "safe_restart: $1 must be a positive integer" >&2; exit 64 ;;
  esac
}

positive_integer DRAIN_POLLS "$DRAIN_POLLS"
positive_integer DRAIN_INTERVAL_SECONDS "$DRAIN_INTERVAL"
positive_integer READY_TIMEOUT_SECONDS "$READY_TIMEOUT"
positive_integer READY_INTERVAL_SECONDS "$READY_INTERVAL"
positive_integer CURL_TIMEOUT_SECONDS "$CURL_TIMEOUT"

health_body() {
  curl -sS --max-time "$CURL_TIMEOUT" "$URL" 2>/dev/null || true
}

is_ready() {
  local body="$1"
  grep -Eq '"status"[[:space:]]*:[[:space:]]*"ok"' <<<"$body" \
    && grep -Eq '"dedicated"[[:space:]]*:[[:space:]]*true' <<<"$body" \
    && grep -Eq '"retail"[[:space:]]*:[[:space:]]*true' <<<"$body" \
    && grep -Eq '"retailBundles"[[:space:]]*:[[:space:]]*true' <<<"$body"
}

for ((i = 1; i <= DRAIN_POLLS; i++)); do
  body=$(health_body)
  clients=$(printf '%s' "$body" | sed -n 's/.*"clients"[[:space:]]*:[[:space:]]*\([0-9][0-9]*\).*/\1/p')
  downloads=$(printf '%s' "$body" | sed -n 's/.*"retailDownloads"[[:space:]]*:[[:space:]]*\([0-9][0-9]*\).*/\1/p')
  if [ -z "$downloads" ]; then
    downloads=$(printf '%s' "$body" | sed -n 's/.*"downloads"[[:space:]]*:[[:space:]]*\([0-9][0-9]*\).*/\1/p')
  fi
  [ -n "$downloads" ] || downloads=0
  retail_busy=0
  grep -Eq '"retailBusy"[[:space:]]*:[[:space:]]*true' <<<"$body" && retail_busy=1
  if [ -z "$clients" ]; then
    echo "safe_restart: healthz unreadable ($body) - poll $i/$DRAIN_POLLS" >&2
  elif [ "$clients" = "0" ] && [ "$downloads" = "0" ] && [ "$retail_busy" = "0" ]; then
    systemctl restart "$SERVICE" \
      || { echo "safe_restart: systemctl restart $SERVICE FAILED" >&2; exit 1; }

    ready_deadline=$((SECONDS + READY_TIMEOUT))
    while :; do
      body=$(health_body)
      if is_ready "$body"; then
        echo "safe_restart: $SERVICE restarted; gateway and dedicated map are ready"
        exit 0
      fi
      ((SECONDS >= ready_deadline)) && break
      remaining=$((ready_deadline - SECONDS))
      delay=$READY_INTERVAL
      ((delay > remaining)) && delay=$remaining
      sleep "$delay"
    done
    echo "safe_restart: $SERVICE restarted but gateway/dedicated readiness failed after ${READY_TIMEOUT}s" >&2
    exit 1
  else
    echo "safe_restart: clients=$clients downloads=$downloads retail_busy=$retail_busy - waiting ($i/$DRAIN_POLLS)"
  fi
  if ((i < DRAIN_POLLS)); then
    sleep "$DRAIN_INTERVAL"
  fi
done

echo "safe_restart: SKIPPED - clients still connected after $DRAIN_POLLS polls" >&2
exit 2
