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

usage() {
  cat <<'EOF'
Usage:
  pi-loop --count N --prompt TEXT [options]
  pi-loop --count N --prompt-file FILE [options]

Options:
  --count N              Number of iterations. Required.
  --prompt TEXT          Base prompt. Required unless --prompt-file is used.
  --prompt-file FILE     Read base prompt from file.
  --cwd DIR              Working dir for pi. Default: current dir.
  --provider NAME        Pass --provider to pi. Default: openai-codex.
  --model MODEL          Pass --model to pi. Default: gpt-5.4-mini.
  --thinking LEVEL       Pass --thinking to pi. Default: low.
  --output-dir DIR       Run artifacts dir. Default: /home/sebas/runtime/pi-loop-runs/TIMESTAMP
  --tail-lines N         Prior output lines injected into next prompt. Default: 120.
  --keep-going           Continue loop even if one iteration fails.
  --no-extensions        Pass --no-extensions to pi. Default: on.
  --allow-extensions     Do not pass --no-extensions.
  -h, --help             Show help.

Examples:
  pi-loop --count 5 --cwd /path/repo --prompt "Implement feature X"
  pi-loop --count 10 --prompt-file ./task.txt --provider openai-codex --model gpt-5.4-mini
EOF
}

COUNT=""
PROMPT=""
PROMPT_FILE=""
CWD="$(pwd)"
PROVIDER="openai-codex"
MODEL="gpt-5.4-mini"
THINKING="low"
TAIL_LINES=120
KEEP_GOING=0
DISABLE_EXTENSIONS=1
STAMP="$(date -u +%Y%m%dT%H%M%SZ)"
OUTPUT_DIR="/home/sebas/runtime/pi-loop-runs/$STAMP"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --count) COUNT="$2"; shift 2 ;;
    --prompt) PROMPT="$2"; shift 2 ;;
    --prompt-file) PROMPT_FILE="$2"; shift 2 ;;
    --cwd) CWD="$2"; shift 2 ;;
    --provider) PROVIDER="$2"; shift 2 ;;
    --model) MODEL="$2"; shift 2 ;;
    --thinking) THINKING="$2"; shift 2 ;;
    --output-dir) OUTPUT_DIR="$2"; shift 2 ;;
    --tail-lines) TAIL_LINES="$2"; shift 2 ;;
    --keep-going) KEEP_GOING=1; shift ;;
    --no-extensions) DISABLE_EXTENSIONS=1; shift ;;
    --allow-extensions) DISABLE_EXTENSIONS=0; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown arg: $1" >&2; usage >&2; exit 1 ;;
  esac
done

if [[ -z "$COUNT" ]]; then
  echo "--count required" >&2
  exit 1
fi
if ! [[ "$COUNT" =~ ^[1-9][0-9]*$ ]]; then
  echo "--count must be a positive integer" >&2
  exit 1
fi
if [[ -n "$PROMPT" && -n "$PROMPT_FILE" ]]; then
  echo "Use either --prompt or --prompt-file, not both" >&2
  exit 1
fi
if [[ -z "$PROMPT" && -z "$PROMPT_FILE" ]]; then
  echo "Prompt required" >&2
  exit 1
fi
if [[ ! -d "$CWD" ]]; then
  echo "cwd not found: $CWD" >&2
  exit 1
fi
if [[ -n "$PROMPT_FILE" ]]; then
  if [[ ! -f "$PROMPT_FILE" ]]; then
    echo "prompt file not found: $PROMPT_FILE" >&2
    exit 1
  fi
  PROMPT="$(cat "$PROMPT_FILE")"
fi
if ! [[ "$TAIL_LINES" =~ ^[0-9]+$ ]]; then
  echo "--tail-lines must be a non-negative integer" >&2
  exit 1
fi

mkdir -p "$OUTPUT_DIR"
CWD_ABS="$(cd "$CWD" && pwd)"
TOOL_USAGE_LOG="$(cd "$(dirname "$0")" && pwd)/tool-usage-log"
LOOP_STARTED_AT_MS="$(python - <<'PY'
import time
print(int(time.time()*1000))
PY
)"
"$TOOL_USAGE_LOG" started loop pi-loop 0 '' '' "$CWD_ABS" '' '' || true
BASE_PROMPT_FILE="$OUTPUT_DIR/base-prompt.txt"
printf '%s\n' "$PROMPT" > "$BASE_PROMPT_FILE"

cat > "$OUTPUT_DIR/run-meta.txt" <<EOF
created_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)
cwd=$CWD_ABS
count=$COUNT
provider=$PROVIDER
model=$MODEL
thinking=$THINKING
tail_lines=$TAIL_LINES
keep_going=$KEEP_GOING
disable_extensions=$DISABLE_EXTENSIONS
base_prompt_file=$BASE_PROMPT_FILE
EOF

build_prompt() {
  local iteration="$1"
  local prev_output_file="$2"

  cat <<EOF
Loop controller note.
You are running inside a repeated multi-session loop.
This is iteration $iteration of $COUNT.
Other iterations may already have changed files in the working tree.
Use the current repo state as source of truth.
Push the task forward as much as possible.
If most core work is already done, verify, refine, fix gaps, tighten docs, or improve the result.
Do not stop early just because prior sessions may have already done the main coding work.
EOF

  if [[ -n "$prev_output_file" && -f "$prev_output_file" && "$TAIL_LINES" -gt 0 ]]; then
    printf '\nPrevious iteration output tail.\n'
    printf '%s\n' '---'
    tail -n "$TAIL_LINES" "$prev_output_file"
    printf '%s\n' '---'
  fi

  printf '\nTask prompt.\n'
  printf '%s\n' '---'
  cat "$BASE_PROMPT_FILE"
  printf '%s\n' '---'
}

PI_ARGS=()
if [[ "$DISABLE_EXTENSIONS" -eq 1 ]]; then
  PI_ARGS+=(--no-extensions)
fi
if [[ -n "$PROVIDER" ]]; then
  PI_ARGS+=(--provider "$PROVIDER")
fi
if [[ -n "$MODEL" ]]; then
  PI_ARGS+=(--model "$MODEL")
fi
if [[ -n "$THINKING" ]]; then
  PI_ARGS+=(--thinking "$THINKING")
fi

LAST_OUTPUT_FILE=""
FAILURES=0

for ((i=1; i<=COUNT; i++)); do
  ITER_DIR="$OUTPUT_DIR/iter-$(printf '%03d' "$i")"
  mkdir -p "$ITER_DIR"

  PROMPT_FILE_ITER="$ITER_DIR/prompt.txt"
  OUTPUT_FILE="$ITER_DIR/output.log"
  STATUS_FILE="$ITER_DIR/status.txt"

  build_prompt "$i" "$LAST_OUTPUT_FILE" > "$PROMPT_FILE_ITER"

  printf 'running\n' > "$STATUS_FILE"
  date -u +"%Y-%m-%dT%H:%M:%SZ" > "$ITER_DIR/started_at.txt"

  set +e
  (
    cd "$CWD_ABS"
    pi "${PI_ARGS[@]}" -p "$(cat "$PROMPT_FILE_ITER")"
  ) 2>&1 | tee "$OUTPUT_FILE"
  EXIT_CODE=${PIPESTATUS[0]}
  set -e

  printf '%s\n' "$EXIT_CODE" > "$ITER_DIR/exit_code.txt"
  date -u +"%Y-%m-%dT%H:%M:%SZ" > "$ITER_DIR/finished_at.txt"

  if [[ "$EXIT_CODE" -eq 0 ]]; then
    printf 'succeeded\n' > "$STATUS_FILE"
  else
    printf 'failed\n' > "$STATUS_FILE"
    FAILURES=$((FAILURES + 1))
    if [[ "$KEEP_GOING" -ne 1 ]]; then
      FINISHED_AT_MS="$(python - <<'PY'
import time
print(int(time.time()*1000))
PY
)"
      "$TOOL_USAGE_LOG" failed loop pi-loop "$((FINISHED_AT_MS - LOOP_STARTED_AT_MS))" "$EXIT_CODE" "iteration $i failed" "$CWD_ABS" '' '' || true
      echo "iteration $i failed with exit code $EXIT_CODE" >&2
      echo "run_dir=$OUTPUT_DIR" >&2
      exit "$EXIT_CODE"
    fi
  fi

  LAST_OUTPUT_FILE="$OUTPUT_FILE"
done

FINISHED_AT_MS="$(python - <<'PY'
import time
print(int(time.time()*1000))
PY
)"
if [[ "$FAILURES" -eq 0 ]]; then
  "$TOOL_USAGE_LOG" succeeded loop pi-loop "$((FINISHED_AT_MS - LOOP_STARTED_AT_MS))" 0 '' "$CWD_ABS" '' '' || true
else
  "$TOOL_USAGE_LOG" failed loop pi-loop "$((FINISHED_AT_MS - LOOP_STARTED_AT_MS))" 1 "$FAILURES iterations failed" "$CWD_ABS" '' '' || true
fi

echo "run_dir=$OUTPUT_DIR"
echo "iterations=$COUNT"
echo "failures=$FAILURES"