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

usage() {
  cat <<'EOF'
Usage: caption_until_complete.sh [YYYY-MM-DD]

Environment overrides:
  CAPTION_ATTEMPT_TIMEOUT   Per-attempt timeout passed to `timeout` (default: 20m)
  CAPTION_SLEEP_SECONDS     Sleep between attempts when incomplete (default: 30)
  CAPTION_MAX_ATTEMPTS      0 = infinite loop until complete (default: 0)
  CAPTION_AI_BIN            Explicit AI command path/binary (default: auto-detect z-ai)
  CAPTION_SKIP_AI           1 = skip AI call (useful for testing)
EOF
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
DATE_ISO="${1:-$(date +%F)}"
DAY_DIR="${ROOT}/PHASES/PHASE-01-INGEST/tweets/archive/${DATE_ISO}"
LOCK="/tmp/twitter-caption-${DATE_ISO}.lock"
LOG_DIR="${ROOT}/PHASES/PHASE-02-CAPTIONS/logs/caption"
ATTEMPT_TIMEOUT="${CAPTION_ATTEMPT_TIMEOUT:-20m}"
SLEEP_SECONDS="${CAPTION_SLEEP_SECONDS:-30}"
MAX_ATTEMPTS="${CAPTION_MAX_ATTEMPTS:-0}"
SKIP_AI="${CAPTION_SKIP_AI:-0}"

MANIFEST="${DAY_DIR}/manifest.json"
WITH_IMAGES="${DAY_DIR}/with_images.csv"
OK_FILE="${DAY_DIR}/images_done.ok"

mkdir -p "${LOG_DIR}" "${DAY_DIR}"

exec 9>"${LOCK}"
flock -n 9 || { echo "[caption] already running date=${DATE_ISO}"; exit 0; }

if [[ ! -f "${MANIFEST}" ]]; then
  echo "[caption] ERROR: missing ${MANIFEST}"
  exit 2
fi

if [[ ! -f "${WITH_IMAGES}" ]]; then
  echo "[caption] ERROR: missing ${WITH_IMAGES}"
  exit 2
fi

if [[ -f "${OK_FILE}" ]]; then
  if python3 "${ROOT}/PHASES/PHASE-02-CAPTIONS/validate_captions.py" "${DATE_ISO}" >/dev/null 2>&1; then
    echo "[caption] already complete (${OK_FILE} present and valid)"
    exit 0
  fi
  echo "[caption] stale ${OK_FILE}; validator failed, removing checkpoint"
  rm -f "${OK_FILE}"
fi

AI_BIN="${CAPTION_AI_BIN:-}"
if [[ -z "${AI_BIN}" && "${SKIP_AI}" != "1" ]]; then
  if command -v z-ai >/dev/null 2>&1; then
    AI_BIN="z-ai"
  elif [[ -x /home/agent/bin/z-ai ]]; then
    AI_BIN="/home/agent/bin/z-ai"
  else
    echo "[caption] ERROR: could not find z-ai. Set CAPTION_AI_BIN=/path/to/z-ai."
    exit 2
  fi
fi

PROMPT="Read ${ROOT}/PHASES/PHASE-02-CAPTIONS/CLAUDE.md and caption missing images for date ${DATE_ISO}. Only append to ${DAY_DIR}/image_descriptions.jsonl."

echo "[caption] start date=${DATE_ISO} time=$(date -Is)"
attempt=0
while true; do
  attempt=$((attempt + 1))
  echo "[caption] attempt=${attempt} time=$(date -Is)"

  ai_rc=0
  if [[ "${SKIP_AI}" == "1" ]]; then
    echo "[caption] CAPTION_SKIP_AI=1; skipping AI step"
  else
    set +e
    timeout "${ATTEMPT_TIMEOUT}" "${AI_BIN}" -p "${PROMPT}"
    ai_rc=$?
    set -e

    if [[ "${ai_rc}" -eq 0 ]]; then
      echo "[caption] AI step finished"
    elif [[ "${ai_rc}" -eq 124 ]]; then
      echo "[caption] AI step timed out after ${ATTEMPT_TIMEOUT}; continuing"
    else
      echo "[caption] AI step exited rc=${ai_rc}; continuing"
    fi
  fi

  python3 "${ROOT}/PHASES/PHASE-02-CAPTIONS/apply_image_descriptions.py" "${DATE_ISO}"

  set +e
  python3 "${ROOT}/PHASES/PHASE-02-CAPTIONS/validate_captions.py" "${DATE_ISO}"
  validate_rc=$?
  set -e

  if [[ "${validate_rc}" -eq 0 ]]; then
    echo "ok $(date -Is)" > "${OK_FILE}"
    echo "[caption] COMPLETE time=$(date -Is)"
    exit 0
  fi

  if [[ "${validate_rc}" -eq 2 ]]; then
    echo "[caption] validator hard error; stopping"
    exit 2
  fi

  if [[ "${MAX_ATTEMPTS}" -gt 0 && "${attempt}" -ge "${MAX_ATTEMPTS}" ]]; then
    echo "[caption] reached CAPTION_MAX_ATTEMPTS=${MAX_ATTEMPTS}; still incomplete"
    exit 1
  fi

  echo "[caption] incomplete; sleeping ${SLEEP_SECONDS}s"
  sleep "${SLEEP_SECONDS}"
done
