#!/usr/bin/env bash
# burn — a token-burning harness around Claude Code or Codex CLI.
# Usage: burn TOKENS [--model MODEL] [--backend claude|codex]
set -euo pipefail

# Resolve this script's directory so we can find sibling files (codex_pricing.py).
SCRIPT_PATH="${BASH_SOURCE[0]}"
while [[ -L "$SCRIPT_PATH" ]]; do
  SCRIPT_DIR="$(cd -P "$(dirname "$SCRIPT_PATH")" && pwd)"
  SCRIPT_PATH="$(readlink "$SCRIPT_PATH")"
  [[ "$SCRIPT_PATH" != /* ]] && SCRIPT_PATH="$SCRIPT_DIR/$SCRIPT_PATH"
done
SCRIPT_DIR="$(cd -P "$(dirname "$SCRIPT_PATH")" && pwd)"

FILLER_TOKENS=1000
# Backend-specific system overhead, measured empirically per CLI.
CLAUDE_SYSTEM_OVERHEAD_TOKENS=9000  # claude -p with --tools ""
CODEX_SYSTEM_OVERHEAD_TOKENS=12000  # codex exec --sandbox read-only
FILLER_LINE="Burn baby burn, this is a disco inferno, burning tokens is for great engineers. "

usage() {
  cat <<EOF
burn — spin up Claude Code or Codex in a loop to burn approximately N tokens.

Usage:
  burn TOKENS [--model MODEL] [--backend claude|codex]

Arguments:
  TOKENS         Target tokens to burn. Minimum per call:
                   claude: $((CLAUDE_SYSTEM_OVERHEAD_TOKENS + FILLER_TOKENS))
                   codex:  $((CODEX_SYSTEM_OVERHEAD_TOKENS + FILLER_TOKENS))

Options:
  --model NAME       Model to pass through (e.g. haiku, sonnet, gpt-5).
  --backend NAME     Which CLI to drive: 'claude' or 'codex'.
                     Default: claude if installed, else codex.
  -h, --help         Show this help and exit.
EOF
}

model=""
tokens=""
backend=""
while [[ $# -gt 0 ]]; do
  case "$1" in
  -h | --help)
    usage
    exit 0
    ;;
  --model)
    if [[ $# -lt 2 ]]; then
      echo "error: --model requires a value" >&2
      exit 2
    fi
    model="$2"
    shift 2
    ;;
  --model=*)
    model="${1#*=}"
    shift
    ;;
  --backend)
    if [[ $# -lt 2 ]]; then
      echo "error: --backend requires a value" >&2
      exit 2
    fi
    backend="$2"
    shift 2
    ;;
  --backend=*)
    backend="${1#*=}"
    shift
    ;;
  --)
    shift
    break
    ;;
  -*)
    echo "error: unknown flag: $1" >&2
    usage >&2
    exit 2
    ;;
  *)
    if [[ -n "$tokens" ]]; then
      echo "error: unexpected extra argument: $1" >&2
      exit 2
    fi
    tokens="$1"
    shift
    ;;
  esac
done

if [[ -z "$tokens" ]]; then
  usage >&2
  exit 2
fi
if ! [[ "$tokens" =~ ^[0-9]+$ ]]; then
  echo "error: TOKENS must be a positive integer (got '$tokens')" >&2
  exit 2
fi

# Resolve backend: explicit flag > auto-detect (claude preferred, codex fallback).
if [[ -z "$backend" ]]; then
  if command -v claude >/dev/null 2>&1; then
    backend=claude
  elif command -v codex >/dev/null 2>&1; then
    backend=codex
  else
    echo "error: neither 'claude' nor 'codex' found on PATH." >&2
    echo "       install Claude Code: https://docs.claude.com/claude-code" >&2
    echo "       or Codex CLI:        https://github.com/openai/codex" >&2
    exit 1
  fi
fi

case "$backend" in
  claude) SYSTEM_OVERHEAD_TOKENS=$CLAUDE_SYSTEM_OVERHEAD_TOKENS ;;
  codex)  SYSTEM_OVERHEAD_TOKENS=$CODEX_SYSTEM_OVERHEAD_TOKENS ;;
  *)
    echo "error: unknown --backend '$backend' (must be 'claude' or 'codex')" >&2
    exit 2
    ;;
esac

if ! command -v "$backend" >/dev/null 2>&1; then
  echo "error: '$backend' CLI not found on PATH." >&2
  case "$backend" in
    claude) echo "       install Claude Code: https://docs.claude.com/claude-code" >&2 ;;
    codex)  echo "       install Codex CLI:   https://github.com/openai/codex" >&2 ;;
  esac
  exit 1
fi
if ! command -v jq >/dev/null 2>&1; then
  echo "error: 'jq' not found on PATH (used to parse CLI JSON output)." >&2
  echo "       install with: brew install jq" >&2
  exit 1
fi

TOKENS_PER_CALL=$((SYSTEM_OVERHEAD_TOKENS + FILLER_TOKENS))
MIN_TOKENS=$TOKENS_PER_CALL
if ((tokens < MIN_TOKENS)); then
  echo "error: tokens must be at least ${MIN_TOKENS} for backend '${backend}' (got ${tokens})" >&2
  exit 2
fi

# ~4 chars/token is the standard rule of thumb for English text.
target_chars=$((FILLER_TOKENS * 4))
filler_len=${#FILLER_LINE}
repeats=$((target_chars / filler_len))
((repeats < 1)) && repeats=1
filler=""
for ((i = 0; i < repeats; i++)); do
  filler+="$FILLER_LINE"
done
prompt=$(printf "Reply with the single word 'burned' and nothing else. Ignore the filler text below; it exists only to consume tokens.\n\n%s" "$filler")

if [[ -t 1 ]]; then
  ORANGE=$'\033[1;33m'
  MAGENTA=$'\033[1;35m'
  CYAN=$'\033[1;36m'
  RED=$'\033[1;31m'
  YELLOW=$'\033[1;93m'
  DIM=$'\033[2m'
  RESET=$'\033[0m'
  HIDE_CURSOR=$'\033[?25l'
  SHOW_CURSOR=$'\033[?25h'
else
  ORANGE=''
  MAGENTA=''
  CYAN=''
  RED=''
  YELLOW=''
  DIM=''
  RESET=''
  HIDE_CURSOR=''
  SHOW_CURSOR=''
fi

PAD='        '

# Four disco frames cycled to make the dancer sway across the floor while
# the ball glitters on the off-beat. Each row is exactly 23 characters wide
# so it sits cleanly inside the booth opening.
disco_frame() {
  case "$1_$2" in
  # frame 0 — centered Travolta: left hand on hip, right arm pointing up-right
  0_0) printf '%s' '           |           ' ;;
  0_1) printf '%s' '          🪩           ' ;;
  0_2) printf '%s' '          /|\          ' ;;
  0_3) printf '%s' '             /         ' ;;
  0_4) printf '%s' '           o/          ' ;;
  0_5) printf '%s' '          <|           ' ;;
  0_6) printf '%s' '           |           ' ;;
  0_7) printf '%s' '          / \          ' ;;
  0_8) printf '%s' '  ▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓  ' ;;
  # frame 1 — lean left, same Travolta pose; sparkles on the chain
  1_0) printf '%s' '      *    |    *      ' ;;
  1_1) printf '%s' '          🪩           ' ;;
  1_2) printf '%s' '          \|/          ' ;;
  1_3) printf '%s' '         /             ' ;;
  1_4) printf '%s' '       o/              ' ;;
  1_5) printf '%s' '      <|               ' ;;
  1_6) printf '%s' '       |               ' ;;
  1_7) printf '%s' '      / \              ' ;;
  1_8) printf '%s' '  ▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒  ' ;;
  # frame 2 — centered mirror: right hand on hip, left arm pointing up-left
  2_0) printf '%s' '           |           ' ;;
  2_1) printf '%s' '          🪩           ' ;;
  2_2) printf '%s' '          /|\          ' ;;
  2_3) printf '%s' '         \             ' ;;
  2_4) printf '%s' '          \o           ' ;;
  2_5) printf '%s' '           |>          ' ;;
  2_6) printf '%s' '           |           ' ;;
  2_7) printf '%s' '          / \          ' ;;
  2_8) printf '%s' '  ▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓  ' ;;
  # frame 3 — lean right, mirror Travolta; sparkles on the chain
  3_0) printf '%s' '    *      |      *    ' ;;
  3_1) printf '%s' '          🪩           ' ;;
  3_2) printf '%s' '          \|/          ' ;;
  3_3) printf '%s' '             \         ' ;;
  3_4) printf '%s' '              \o       ' ;;
  3_5) printf '%s' '               |>      ' ;;
  3_6) printf '%s' '               |       ' ;;
  3_7) printf '%s' '              / \      ' ;;
  3_8) printf '%s' '  ▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒▓▒  ' ;;
  esac
}

disco_color_for_row() {
  case "$1" in
  0 | 2) printf '%s' "$ORANGE" ;;  # chain sparkles + light rays
  1 | 8) printf '%s' "$MAGENTA" ;; # disco ball + dance floor
  *) printf '%s' "$CYAN" ;;        # dancer (and blank rows)
  esac
}

# Render `count` copies of `char` colored across a cycling fire palette
# (red → orange → yellow). `offset` rotates the cycle so the strip flickers.
fire_strip() {
  local count=$1 char=$2 offset=${3:-0}
  local palette=("$RED" "$ORANGE" "$YELLOW")
  local n=${#palette[@]}
  local i out=""
  for ((i = 0; i < count; i++)); do
    out+="${palette[$(((i + offset) % n))]}${char}"
  done
  printf '%s' "$out"
}

print_disco_row() {
  local frame=$1 row=$2 content color fl
  content=$(disco_frame "$frame" "$row")
  color=$(disco_color_for_row "$row")
  fl=$(fire_strip 3 '░' "$frame")
  printf '%s%s│%s%s│%s%s%s%s%s│%s%s│%s\n' \
    "$PAD" "$RED" "$fl" "$YELLOW" "$RESET" \
    "$color" "$content" "$RESET" \
    "$YELLOW" "$fl" "$RED" "$RESET"
}

draw_disco() {
  local frame=${1:-0}
  local flames fr3
  flames=$(fire_strip 31 '▲')
  fr3=$(fire_strip 3 '░')
  printf '%s%s┌───────────────────────────────┐%s\n' "$PAD" "$RED" "$RESET"
  printf '%s%s│%s%s│%s\n' "$PAD" "$RED" "$flames" "$RED" "$RESET"
  printf '%s%s│%s%s┌───────────────────────┐%s%s│%s\n' \
    "$PAD" "$RED" "$fr3" "$YELLOW" "$fr3" "$RED" "$RESET"
  for row in 0 1 2 3 4 5 6 7 8; do
    print_disco_row "$frame" "$row"
  done
  printf '%s%s│%s%s└───────────────────────┘%s%s│%s\n' \
    "$PAD" "$RED" "$fr3" "$YELLOW" "$fr3" "$RED" "$RESET"
  printf '%s%s└───────────────────────────────┘%s\n' "$PAD" "$RED" "$RESET"
}

# Distance (in lines) from the progress-bar cursor up to the top inner row:
# blank line above bar + outer-bottom + inner-bottom + 9 inner rows = 12.
DISCO_OFFSET=12

redraw_disco() {
  local frame=$1
  printf '\033[s'                     # save cursor
  printf '\033[%dA\r' "$DISCO_OFFSET" # move up to top inner row
  for row in 0 1 2 3 4 5 6 7 8; do
    print_disco_row "$frame" "$row"
  done
  printf '\033[u' # restore cursor
}

burned=0
iteration=0
total_cost=0.000000
BAR_WIDTH=30
cli_pid=""
last_output=""

# Build CLI command for the chosen backend.
case "$backend" in
  claude)
    cli_cmd=(claude -p --output-format json --no-session-persistence --tools "")
    [[ -n "$model" ]] && cli_cmd+=(--model "$model")
    # --tools is variadic, so terminate flag parsing before the prompt.
    cli_cmd+=(--)
    ;;
  codex)
    cli_cmd=(codex exec --json --skip-git-repo-check --sandbox read-only)
    [[ -n "$model" ]] && cli_cmd+=(--model "$model")
    ;;
esac

# Resolve whether we can report dollar cost for this run:
#   - claude reports it directly in JSON.
#   - codex doesn't, so we estimate from token counts using codex_rates.sh.
#     If the rates file is missing or the chosen model isn't in the table,
#     we silently fall back to tokens-only display.
cost_reported=0
cost_model=""
if [[ "$backend" == claude ]]; then
  cost_reported=1
elif [[ "$backend" == codex ]]; then
  if [[ -f "$SCRIPT_DIR/codex_rates.sh" ]]; then
    # shellcheck source=codex_rates.sh
    source "$SCRIPT_DIR/codex_rates.sh"
    # Assume gpt-5.5 unless the user passed --model.
    cost_model="${model:-gpt-5.5}"
    if codex_has_rates "$cost_model"; then
      cost_reported=1
    fi
  fi
fi

format_cost_field() {
  if ((cost_reported)); then
    printf '$%s' "$total_cost"
  else
    printf 'cost n/a'
  fi
}

draw_bar() {
  local pct filled empty bar="" i suffix="${1:-}"
  pct=$((burned * 100 / tokens))
  ((pct > 100)) && pct=100
  filled=$((pct * BAR_WIDTH / 100))
  empty=$((BAR_WIDTH - filled))
  for ((i = 0; i < filled; i++)); do bar+='█'; done
  for ((i = 0; i < empty; i++)); do bar+='░'; done
  # \033[K clears any leftover characters from a longer previous draw.
  printf "\r[%s] %3d%% • %d/%d tokens • %s • %d call(s)%s\033[K" \
    "$bar" "$pct" "$burned" "$tokens" "$(format_cost_field)" "$iteration" "$suffix"
}

restore_cursor() {
  printf '%s' "$SHOW_CURSOR"
}
trap restore_cursor EXIT

CELEBRATIONS=(
  "🔧 You are an true AI-native engineer!"
  "🪄 You are an AI wizard!"
  "🕺 Tokens burned. Groove restored."
  "✨ You are a master of vibe-coding."
  "🚀 You shipped faster than the flame could spread."
  "💃 The code is hot and the vibes are immaculate."
  "🪩 Groovy job. I've never seen someone so productive!"
  "🏆 Another glorious victory for human-machine collaboration."
  "⚡ You turned tokens into pure developer energy."
  "🎩 Behold: a masterclass in responsible token combustion."
  "🛠️  Great work. The token furnace is pleased."
)

# Round a dollar amount to the nearest cent.
round_cents() {
  awk -v c="$1" 'BEGIN{printf "%.2f", c}'
}

# Format a non-negative integer with thousands-grouping commas (21826 -> 21,826).
fmt_num() {
  local n=$1 result=""
  while ((n >= 1000)); do
    result=$(printf ",%03d%s" $((n % 1000)) "$result")
    n=$((n / 1000))
  done
  printf "%d%s" "$n" "$result"
}

on_interrupt() {
  if [[ -n "$cli_pid" ]]; then
    kill "$cli_pid" 2>/dev/null || true
    wait "$cli_pid" 2>/dev/null || true
  fi
  local burned_fmt cost_msg
  burned_fmt=$(fmt_num "$burned")
  if ((cost_reported)); then
    cost_msg=$(printf "💰 spent \$%s" "$(round_cents "$total_cost")")
  else
    cost_msg="💰 cost not reported by codex"
  fi
  printf "\n\n🛑 stopped early.\n\n"
  printf "   🔥 torched %s tokens across %s call(s) • %s\n\n" "$burned_fmt" "$iteration" "$cost_msg"
  exit 130
}
trap on_interrupt INT

# Run one CLI call in the background while animating flames. On success
# leaves the parsed stdout in $last_output. Returns the CLI exit status.
run_cli() {
  local output_file
  output_file=$(mktemp)
  # </dev/null prevents codex from waiting on stdin when run in background.
  "${cli_cmd[@]}" "$prompt" >"$output_file" 2>/dev/null </dev/null &
  cli_pid=$!

  local frame=0
  while kill -0 "$cli_pid" 2>/dev/null; do
    redraw_disco $((frame % 4))
    frame=$((frame + 1))
    sleep 0.4
  done

  if wait "$cli_pid" 2>/dev/null; then
    last_output=$(<"$output_file")
    rm -f "$output_file"
    cli_pid=""
    return 0
  else
    rm -f "$output_file"
    cli_pid=""
    return 1
  fi
}

# Parse one CLI call's output into "<tokens> <cost>" (tab-separated).
parse_usage() {
  case "$backend" in
    claude)
      jq -r '
        [ ((.usage.input_tokens // 0) +
           (.usage.cache_creation_input_tokens // 0) +
           (.usage.cache_read_input_tokens // 0) +
           (.usage.output_tokens // 0)),
          (.total_cost_usd // 0)
        ] | @tsv' <<<"$last_output"
      ;;
    codex)
      # codex emits JSONL; sum usage across any turn.completed events.
      # input_tokens already includes cached_input_tokens, so don't double-count
      # for the token total. Cost is computed locally from per-model rates.
      local inp cached out cost
      read -r inp cached out < <(jq -s -r '
        [ .[] | select(.type? == "turn.completed") | .usage ]
        | [ ([.[] | .input_tokens // 0]        | add // 0),
            ([.[] | .cached_input_tokens // 0] | add // 0),
            ([.[] | .output_tokens // 0]       | add // 0) ]
        | @tsv' <<<"$last_output")
      if ((cost_reported)); then
        cost=$(codex_cost "$cost_model" "$inp" "$cached" "$out")
      else
        cost=0
      fi
      printf '%d\t%s\n' "$((inp + out))" "$cost"
      ;;
  esac
}

printf '%s' "$HIDE_CURSOR"
echo
draw_disco 0
echo
draw_bar

while ((burned < tokens)); do
  iteration=$((iteration + 1))

  if ! run_cli; then
    draw_bar " (retrying)"
    if ! run_cli; then
      printf "\nerror: two consecutive %s calls failed; aborting.\n" "$backend" >&2
      exit 1
    fi
  fi

  read -r this_call this_cost < <(parse_usage)

  if ! [[ "$this_call" =~ ^[0-9]+$ ]]; then
    printf "\nerror: could not parse token usage from %s output\n" "$backend" >&2
    exit 1
  fi

  burned=$((burned + this_call))
  total_cost=$(awk -v a="$total_cost" -v b="$this_cost" 'BEGIN{printf "%.6f", a+b}')
  draw_bar
done

msg="${CELEBRATIONS[RANDOM % ${#CELEBRATIONS[@]}]}"
burned_fmt=$(fmt_num "$burned")
tokens_fmt=$(fmt_num "$tokens")

if ((cost_reported)); then
  cost_rounded=$(round_cents "$total_cost")
  printf "\n\n 🔥 torched %s tokens across %s call(s) (target: %s) • 💰 spent \$%s 💵💵💵\n\n" \
    "$burned_fmt" "$iteration" "$tokens_fmt" "$cost_rounded"
else
  printf "\n\n 🔥 torched %s tokens across %s call(s) (target: %s) • 💰 cost not reported by codex\n\n" \
    "$burned_fmt" "$iteration" "$tokens_fmt"
fi
printf " %s\n\n" "$msg"