#!/usr/bin/env bash
# Training runner: compile PPO_Bot, then loop the Java battle runner with crash recovery.
#
# Usage:
#   ./run.sh [opponent] [rounds]
#
# Arguments (also settable via env vars):
#   opponent  — bot name under SAMPLE_BOTS_DIR (default: Target)
#   rounds    — total rounds to run (default: 100)
#
# Env vars:
#   TANK_ROYALE_JAR   — path to robocode-tankroyale-runner.jar
#   SAMPLE_BOTS_DIR   — path to sample bots archive
#   PPOB_LOG_FILE     — JSON lines log file (default: training_log.jsonl in repo root)
#   PPOB_LR, PPOB_CLIP_EPSILON, PPOB_ENTROPY_COEFF, PPOB_VALUE_LOSS_COEFF,
#   PPOB_MAX_GRAD_NORM, PPOB_GAMMA, PPOB_LAM, PPOB_EPOCHS, PPOB_MINI_BATCH_SIZE
#     — hyperparameters (all optional, PPO_Bot uses sensible defaults)
set -uo pipefail

ENV_FILE="$(cd "$(dirname "$0")" && pwd)/training.env"
[ -f "$ENV_FILE" ] && set -a && . "$ENV_FILE" && set +a

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

OPPONENT="${1:-${TRAINING_OPPONENT:-Target}}"
TOTAL_ROUNDS="${2:-${TRAINING_ROUNDS:-100}}"

JAR="${TANK_ROYALE_JAR:-/home/davide/Projects/tank-royale/runner/examples/lib/robocode-tankroyale-runner.jar}"
PPO_BOT_SRC="$REPO_ROOT/PPO_Bot"

export PPO_BOT_DIR="$PPO_BOT_SRC"
export SAMPLE_BOTS_DIR="${SAMPLE_BOTS_DIR:-/home/davide/Projects/tank-royale/sample-bots/java/build/archive}"
export PPOB_LOG_FILE="${PPOB_LOG_FILE:-$REPO_ROOT/training_log.jsonl}"

echo "=== Training runner ==="
echo "Opponent:     $OPPONENT"
echo "Total rounds: $TOTAL_ROUNDS"
echo "Log file:     $PPOB_LOG_FILE"
echo ""

# ── 1. Compile PPO_Bot ────────────────────────────────────────────────────────
echo ">>> Compiling PPO_Bot..."
(cd "$PPO_BOT_SRC" && nimble build -d:release)
echo ">>> PPO_Bot compiled."

# ── 2. Compile RunTraining.java ───────────────────────────────────────────────
cd "$SCRIPT_DIR"
echo ">>> Compiling RunTraining.java..."
javac -cp "$JAR" RunTraining.java
echo ">>> RunTraining compiled."

# ── 3. Run with crash-restart loop ───────────────────────────────────────────
# We pass total rounds to the Java runner. On crash, restart from where we left
# off by reading the round counter PPO_Bot persists to weights/round_counter.txt.
# ponytail: simple restart from persisted counter; no partial-round recovery needed

ROUND_COUNTER_FILE="$PPO_BOT_SRC/weights/round_counter.txt"

remaining_rounds() {
  local done=0
  if [ -f "$ROUND_COUNTER_FILE" ]; then
    done=$(cat "$ROUND_COUNTER_FILE" 2>/dev/null || echo 0)
  fi
  echo $(( TOTAL_ROUNDS - done ))
}

attempts=0
while true; do
  remaining=$(remaining_rounds)
  if [ "$remaining" -le 0 ]; then
    echo ">>> All $TOTAL_ROUNDS rounds complete."
    break
  fi

  attempts=$(( attempts + 1 ))
  if [ "$attempts" -gt 1 ]; then
    echo ">>> Restart #$attempts — $remaining rounds remaining..."
    sleep 2
  fi

  # ponytail: cap each battle at 60 rounds — matches the certified config.
  # Round-end reward is cumulative totalScore/50; long battles grow it unboundedly
  # and blow up the critic's value scale (seen at round 3160, vLoss 25745 by 3871).
  # Short battles bound the reward, and fresh-process restarts clear thread state.
  battle_rounds=$(( remaining > 60 ? 60 : remaining ))
  echo ">>> Running $battle_rounds rounds (opponent: $OPPONENT)..."
  java -cp ".:$JAR" RunTraining "$OPPONENT" "$battle_rounds" && break || {
    exit_code=$?
    echo ">>> Java runner exited with code $exit_code — checking if done..."
    remaining=$(remaining_rounds)
    if [ "$remaining" -le 0 ]; then
      echo ">>> All rounds complete despite non-zero exit."
      break
    fi
    echo ">>> Crash detected ($remaining rounds remain). Restarting..."
  }
done

echo ">>> Training done. Log at $PPOB_LOG_FILE"