#!/usr/bin/env bash
#
# Ralph Wiggum Loop — Autonomous agent iteration
#
# Based on Geoffrey Huntley's approach:
# - Each iteration spawns a FRESH agent with clean context
# - Agent reads the plan, picks ONE task, implements, tests, commits, exits
# - Loop restarts until all tasks are done
#
# Session limit handling:
# - Detects Claude Pro usage limit messages in agent output
# - Polls every SESSION_POLL_INTERVAL seconds until the session resets
# - Resumes the same iteration automatically — no manual intervention needed
#
# Usage:
#   ./ralph-loop.sh              # Build mode (default)
#   ./ralph-loop.sh plan         # Planning mode (create IMPLEMENTATION_PLAN.md)
#   ./ralph-loop.sh --max 20     # Limit to 20 iterations
#   ./ralph-loop.sh --agent claude  # Use claude (default)
#   ./ralph-loop.sh --agent codex   # Use OpenAI Codex CLI
#   ./ralph-loop.sh --agent aider   # Use Aider
#   ./ralph-loop.sh --agent gemini  # Use Gemini CLI
#   ./ralph-loop.sh --agent custom  # Use custom agent (see below)
#
set -euo pipefail

MODE="${1:-build}"
MAX_ITERATIONS=50
AGENT="claude"
PLAN_FILE="IMPLEMENTATION_PLAN.md"
SPEC_FILE="PROJECT-SPEC.md"
AGENT_FILE="AGENT.md"
LOG_DIR=".ralph-logs"

# How often (in seconds) to probe whether the session has reset.
# Default: 10 minutes. Adjust down if you want faster recovery.
SESSION_POLL_INTERVAL="${SESSION_POLL_INTERVAL:-600}"

# Parse arguments
shift 2>/dev/null || true
while [[ $# -gt 0 ]]; do
  case "$1" in
    --max) MAX_ITERATIONS="$2"; shift 2 ;;
    --agent) AGENT="$2"; shift 2 ;;
    *) echo "Unknown option: $1"; exit 1 ;;
  esac
done

mkdir -p "$LOG_DIR"

# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

log()     { echo -e "${BLUE}[ralph]${NC} $1"; }
success() { echo -e "${GREEN}[ralph]${NC} $1"; }
warn()    { echo -e "${YELLOW}[ralph]${NC} $1"; }
error()   { echo -e "${RED}[ralph]${NC} $1"; }
info()    { echo -e "${CYAN}[ralph]${NC} $1"; }

# Check prerequisites
if [[ ! -f "$SPEC_FILE" ]]; then
  error "Missing $SPEC_FILE — create your project spec first."
  exit 1
fi

if [[ ! -f "$AGENT_FILE" ]]; then
  warn "No $AGENT_FILE found. Using default agent instructions."
fi

run_agent() {
  local iteration=$1
  local mode=$2
  local logfile="$LOG_DIR/iteration-${iteration}.log"
  local prompt=""

  if [[ "$mode" == "plan" ]]; then
    prompt="Read PROJECT-SPEC.md. Decompose the project into discrete, testable tasks ordered by dependency. Write the plan to IMPLEMENTATION_PLAN.md with checkboxes. Output <promise>PLANNED</promise> when done."
  else
    prompt="Read AGENT.md (if it exists) for your instructions. Follow the core loop: orient, pick one task, implement, verify, commit, exit."
  fi

  log "Iteration $iteration ($mode mode) — starting fresh agent..."

  # Disable pipefail around the agent call so a non-zero claude exit doesn't
  # kill the script. We inspect the log content instead.
  set +e
  case "$AGENT" in
    claude)
      echo "$prompt" | claude -p --output-format text 2>&1 | tee "$logfile"
      ;;
    codex)
      echo "$prompt" | codex 2>&1 | tee "$logfile"
      ;;
    aider)
      aider --message "$prompt" --yes 2>&1 | tee "$logfile"
      ;;
    gemini)
      echo "$prompt" | gemini-cli 2>&1 | tee "$logfile"
      ;;
    custom)
      if [[ -x "./custom-agent.sh" ]]; then
        ./custom-agent.sh "$prompt" 2>&1 | tee "$logfile"
      else
        error "Custom agent selected but ./custom-agent.sh not found or not executable"
        exit 1
      fi
      ;;
    *)
      error "Unknown agent: $AGENT"
      error "Supported agents: claude, codex, aider, gemini, custom"
      exit 1
      ;;
  esac
  set -e

  return 0
}

# Probe whether claude is available by sending a trivial request.
# Returns 0 if available, 1 if still rate-limited or erroring.
probe_session() {
  local probe_log="$LOG_DIR/probe.log"
  set +e
  echo "Reply with the single word OK and nothing else." \
    | claude -p --output-format text > "$probe_log" 2>&1
  local rc=$?
  set -e

  if [[ $rc -ne 0 ]]; then
    return 1
  fi
  # Also check the output doesn't contain a limit message
  if grep -qi 'usage limit\|rate limit\|limit reached\|exceeded.*limit' "$probe_log" 2>/dev/null; then
    return 1
  fi
  return 0
}

check_output() {
  local logfile="$1"

  # Session / usage limit — must check BEFORE generic promise checks
  if grep -qi 'usage limit\|rate limit\|limit reached\|exceeded.*limit\|Claude AI usage' "$logfile" 2>/dev/null; then
    return 4  # Rate limited
  fi

  if grep -q '<promise>DONE</promise>' "$logfile" 2>/dev/null; then
    return 0  # Done
  elif grep -q '<promise>STUCK</promise>' "$logfile" 2>/dev/null; then
    return 2  # Stuck — needs human intervention
  elif grep -q '<promise>ERROR</promise>' "$logfile" 2>/dev/null; then
    return 3  # Unrecoverable error
  else
    return 1  # Normal iteration — continue
  fi
}

wait_for_session_reset() {
  local iteration=$1
  warn "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
  warn "Session usage limit hit during iteration $iteration."
  warn "Will probe every ${SESSION_POLL_INTERVAL}s until session resets."
  warn "No manual action needed — loop will resume automatically."
  warn "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

  local attempt=0
  while true; do
    ((attempt++))
    local next_check
    next_check=$(date -d "+${SESSION_POLL_INTERVAL} seconds" '+%H:%M:%S' 2>/dev/null \
      || date -v "+${SESSION_POLL_INTERVAL}S" '+%H:%M:%S' 2>/dev/null \
      || echo "soon")
    info "Probe attempt $attempt — next check at $next_check..."
    sleep "$SESSION_POLL_INTERVAL"

    if probe_session; then
      success "Session available! Resuming iteration $iteration..."
      return 0
    else
      warn "Still rate-limited (attempt $attempt). Waiting another ${SESSION_POLL_INTERVAL}s..."
    fi
  done
}

# ─── Main ────────────────────────────────────────────────────────────────────

if [[ "$MODE" == "plan" ]]; then
  log "Planning mode — creating implementation plan..."
  run_agent 0 plan
  success "Plan created. Review $PLAN_FILE, then run: ./ralph-loop.sh"
  exit 0
fi

log "Starting Ralph Wiggum loop (max $MAX_ITERATIONS iterations)"
log "Agent:               $AGENT"
log "Spec:                $SPEC_FILE"
log "Plan:                $PLAN_FILE"
log "Poll interval:       ${SESSION_POLL_INTERVAL}s (session limit recovery)"
echo ""

i=1
while [[ $i -le $MAX_ITERATIONS ]]; do
  run_agent "$i" build
  logfile="$LOG_DIR/iteration-${i}.log"

  # Capture return value without triggering set -e
  check_output "$logfile" || status=$?
  status=${status:-0}

  case $status in
    0)
      success "ALL TASKS COMPLETE after $i iterations!"
      exit 0
      ;;
    2)
      warn "Agent is stuck on iteration $i. Review $logfile and intervene."
      exit 1
      ;;
    3)
      error "Agent encountered an error on iteration $i. Review $logfile."
      exit 1
      ;;
    4)
      # Rate limited — wait for reset, then retry the SAME iteration
      wait_for_session_reset "$i"
      # Do NOT increment i — retry the same task
      ;;
    1)
      log "Iteration $i complete. Restarting with fresh context..."
      echo ""
      sleep 2
      ((i++))
      ;;
  esac
done

warn "Reached max iterations ($MAX_ITERATIONS). Review progress in $PLAN_FILE."
exit 1