#!/bin/zsh
# api-retry — re-run a command with exponential backoff + API liveness probe.
#
# Why this exists:
#   Claude Code / any Anthropic SDK call can fail with "Stream idle timeout —
#   partial response received" when the server pauses streaming mid-response,
#   or with transient network errors. Bare `until cmd; do sleep N; done`
#   retries blindly even when the network is still down, wasting tokens
#   and obscuring the real cause. This wrapper:
#     1. runs the command
#     2. if it fails, waits with exponential backoff
#     3. then HEAD-probes the API endpoint until it actually answers
#     4. only then retries
#   Repeat until success or --max-attempts reached.
#
# Usage:
#   api-retry [options] -- <cmd> [args...]
#
# Options:
#   --probe URL              probe URL (default: api.anthropic.com)
#   --probe-timeout SEC      per-probe timeout (default: 3)
#   --probe-interval SEC     wait between probes when down (default: 15)
#   --max-attempts N         0 = forever (default: 0)
#   --initial-backoff SEC    first backoff (default: 5)
#   --backoff-cap SEC        max backoff (default: 300)
#   --quiet                  no banner / status lines
#
# Examples:
#   api-retry -- claude --print "summarize foo.md"
#   api-retry --max-attempts 10 -- bash my-script.sh
#   api-retry --probe https://other.api/healthz -- curl https://other.api/x
#
# Note on interactive Claude Code sessions:
#   This wraps single-shot commands. If your already-running interactive
#   Claude Code session hits "Stream idle timeout — partial response", the
#   wrapper cannot resume that conversation from the outside. For interactive
#   sessions, prefer raising ANTHROPIC_API_TIMEOUT_MS=600000 instead.

set -u

usage() {
  awk 'NR == 1 { next } /^set -/ { exit } { sub(/^# ?/, ""); print }' "${(%):-%x}"
}

PROBE='https://api.anthropic.com/v1/messages'
PROBE_TIMEOUT=3
PROBE_INTERVAL=15
MAX_ATTEMPTS=0
BACKOFF=5
BACKOFF_CAP=300
QUIET=0

c_grn() { [[ $QUIET -eq 1 ]] || printf '\033[32m[api-retry] %s\033[0m\n' "$*" >&2; }
c_ylw() { [[ $QUIET -eq 1 ]] || printf '\033[33m[api-retry] %s\033[0m\n' "$*" >&2; }
c_red() { printf '\033[31m[api-retry] %s\033[0m\n' "$*" >&2; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    --probe)            PROBE="$2";           shift 2;;
    --probe-timeout)    PROBE_TIMEOUT="$2";   shift 2;;
    --probe-interval)   PROBE_INTERVAL="$2";  shift 2;;
    --max-attempts)     MAX_ATTEMPTS="$2";    shift 2;;
    --initial-backoff)  BACKOFF="$2";         shift 2;;
    --backoff-cap)      BACKOFF_CAP="$2";     shift 2;;
    --quiet)            QUIET=1;              shift;;
    --) shift; break;;
    -h|--help) usage; exit 0;;
    *) echo "api-retry: unknown option: $1" >&2; exit 2;;
  esac
done

if [[ $# -eq 0 ]]; then
  echo "api-retry: no command given (use -- before the command)" >&2
  exit 2
fi

probe_until_alive() {
  # 200/401/403/405 are all "server is up and routing"; only DNS / TCP / TLS /
  # idle-cutoff errors register as 000, which is the only state we wait for.
  # curl with -w always emits a 3-digit code (000 on connect failure) so we
  # don't need an `|| echo` fallback that would concatenate strings.
  while :; do
    code=$(curl -s -o /dev/null -w '%{http_code}' \
                -m "$PROBE_TIMEOUT" -X HEAD "$PROBE" 2>/dev/null)
    code=${code:-000}
    if [[ "$code" != '000' ]]; then
      c_grn "probe ok ($PROBE → HTTP $code)"
      return 0
    fi
    c_ylw "$PROBE unreachable; sleep ${PROBE_INTERVAL}s"
    sleep "$PROBE_INTERVAL"
  done
}

attempt=1
while :; do
  if [[ $attempt -gt 1 ]]; then
    c_ylw "attempt $attempt of ${MAX_ATTEMPTS:-∞}"
  fi

  # NOTE: capture rc immediately; `if cmd; then …; fi` masks it ($? becomes 0
  # for the if itself in bash/zsh).
  "$@"
  rc=$?
  if [[ $rc -eq 0 ]]; then
    [[ $attempt -gt 1 ]] && c_grn "succeeded on attempt $attempt"
    exit 0
  fi

  if [[ $MAX_ATTEMPTS -gt 0 && $attempt -ge $MAX_ATTEMPTS ]]; then
    c_red "giving up after $attempt attempts (last rc=$rc)"
    exit "$rc"
  fi

  c_ylw "attempt $attempt failed (rc=$rc); backoff ${BACKOFF}s, then probe"
  sleep "$BACKOFF"
  probe_until_alive

  attempt=$((attempt + 1))
  BACKOFF=$(( BACKOFF * 2 ))
  [[ $BACKOFF -gt $BACKOFF_CAP ]] && BACKOFF=$BACKOFF_CAP
done
