#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────────────
# ab_run.sh — run an A/B session of N arms vs real DrussGT.
#
#   tools/ab/ab_run.sh --arms arms.txt --runs 7 --outdir /tmp/ab/foo --conc 7
#
# What it does, once per SESSION (not once per arm):
#   1. builds ONE frozen ModularBot from the CURRENT HEAD via `git archive HEAD`
#      (so a dirty tree cannot leak into the measurement) and records the commit
#      sha + binary sha256;
#   2. for every arm × run, launches a real battle against DrussGT in parallel,
#      each with its own ports (the runner picks ephemeral ones), its own
#      DrussGT botdir/data, its own ModularBot botdir and its own output files.
#
# Arm file format (one arm per line):
#     name | ENV_VAR=value ENV_VAR2=value2 | optional label
#   `name` alone or `name |` are also accepted (no arm env). Blank lines and
#   lines starting with `#` are ignored. See arms.example.txt.
#
# Output layout (what ab_analyze.py reads):
#     <outdir>/session.json                   commit, binary sha, arms, runs…
#     <outdir>/<arm>/run<N>.jsonl             tick capture
#     <outdir>/<arm>/run<N>.jsonl.rounds.json round tick boundaries
#     <outdir>/<arm>/run<N>.jsonl.results.json
#     <outdir>/<arm>/run<N>.events.jsonl      fire/hit/death sidecar
#     <outdir>/<arm>/run<N>.battle.log        capture stdout (firstPlaces, …)
#     <outdir>/<arm>/run<N>.bot.stdout.log    ModularBot boot env report (liveness)
#
# Cleanup: every job runs in its own process group (setsid); on EXIT/INT/TERM
# this script SIGTERMs/SIGKILLs those groups and, as a backstop, pkills anything
# whose command line references this session's outdir. So a Ctrl-C does not
# leave orphan battles running.
# ─────────────────────────────────────────────────────────────────────────────
set -euo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
RUN_BRIDGE="$REPO/tools/robocode_shim/run_bridge_battle.sh"
MAKE_BOTDIR="$REPO/tools/robocode_shim/make_botdir.sh"

ARMS_FILE=""
RUNS=7
CONC=7
OUTDIR=""
ROUNDS=7

usage() {
  sed -n '2,40p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//'
  exit "${1:-0}"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --arms)   ARMS_FILE="$2"; shift 2;;
    --runs)   RUNS="$2"; shift 2;;
    --outdir) OUTDIR="$2"; shift 2;;
    --conc)   CONC="$2"; shift 2;;
    --rounds) ROUNDS="$2"; shift 2;;
    -h|--help) usage 0;;
    *) echo "unknown argument: $1" >&2; usage 1;;
  esac
done

[[ -n "$ARMS_FILE" ]] || { echo "ERROR: --arms <file> is required" >&2; usage 1; }
[[ -f "$ARMS_FILE" ]] || { echo "ERROR: arm file not found: $ARMS_FILE" >&2; exit 1; }
[[ -n "$OUTDIR" ]]    || { echo "ERROR: --outdir <dir> is required" >&2; usage 1; }
[[ "$OUTDIR" != "/" && "$OUTDIR" != "" ]] || { echo "ERROR: refusing outdir '$OUTDIR'" >&2; exit 1; }
OUTDIR="$(mkdir -p "$OUTDIR" && cd "$OUTDIR" && pwd)"

# ── prerequisites ────────────────────────────────────────────────────────────
ROBOCODE_JAR="${ROBOCODE_JAR:-/tmp/robocode/install/libs/robocode.jar}"
DRUSSGT_JAR="${DRUSSGT_JAR:-/tmp/drussgt/DrussGT.jar}"
RUNNER_JAR="${TR_RUNNER_JAR:-/home/davide/Projects/tank-royale/runner/examples/lib/robocode-tankroyale-runner.jar}"
BOT_API_JAR="${TR_BOT_API_JAR:-$HOME/Downloads/sample-bots-java-1.0.2/lib/robocode-tankroyale-bot-api-1.0.2.jar}"
DRUSSGT_SHARED="${TR_DRUSSGT_BOTDIR:-/tmp/tr_bots/DrussGT}"

missing=0
for f in "$ROBOCODE_JAR" "$DRUSSGT_JAR" "$RUNNER_JAR" "$BOT_API_JAR"; do
  if [[ ! -f "$f" ]]; then echo "ERROR: missing prerequisite: $f" >&2; missing=1; fi
done
command -v nim >/dev/null 2>&1 || { echo "ERROR: nim not on PATH" >&2; missing=1; }
command -v git >/dev/null 2>&1 || { echo "ERROR: git not on PATH" >&2; missing=1; }
[[ -f "$RUN_BRIDGE" ]]  || { echo "ERROR: missing $RUN_BRIDGE" >&2; missing=1; }
[[ -f "$MAKE_BOTDIR" ]] || { echo "ERROR: missing $MAKE_BOTDIR" >&2; missing=1; }
(( missing == 0 )) || { echo "ERROR: prerequisites missing — not starting a broken session" >&2; exit 1; }

if [[ ! -d "$DRUSSGT_SHARED" ]]; then
  echo "[ab] recreating shared DrussGT botdir: $DRUSSGT_SHARED"
  "$MAKE_BOTDIR" "$DRUSSGT_SHARED" >/dev/null
fi

# ── parse the arm file into parallel arrays ──────────────────────────────────
trim() { local s="$1"; s="${s#"${s%%[![:space:]]*}"}"; s="${s%"${s##*[![:space:]]}"}"; printf '%s' "$s"; }

ARM_NAMES=(); ARM_ENVS=(); ARM_LABELS=()
while IFS= read -r line || [[ -n "$line" ]]; do
  [[ -z "${line//[[:space:]]/}" ]] && continue
  [[ "$line" =~ ^[[:space:]]*# ]] && continue
  name="$(trim "${line%%|*}")"
  rest="${line#*|}"
  if [[ "$line" == *"|"* ]]; then
    envspec="$(trim "${rest%%|*}")"
    label="$(trim "${rest#*|}")"
  else
    envspec=""; label=""
  fi
  [[ -n "$name" ]] || { echo "ERROR: arm with empty name in $ARMS_FILE" >&2; exit 1; }
  ARM_NAMES+=("$name"); ARM_ENVS+=("$envspec"); ARM_LABELS+=("$label")
done < "$ARMS_FILE"
(( ${#ARM_NAMES[@]} > 0 )) || { echo "ERROR: no arms parsed from $ARMS_FILE" >&2; exit 1; }
# duplicate names would silently overwrite each other's run files
dupes="$(printf '%s\n' "${ARM_NAMES[@]}" | sort | uniq -d)"
[[ -z "$dupes" ]] || { echo "ERROR: duplicate arm name(s): $dupes" >&2; exit 1; }

# ── fresh session dir ────────────────────────────────────────────────────────
rm -rf "$OUTDIR"
mkdir -p "$OUTDIR"
WORK="$OUTDIR/.work"; mkdir -p "$WORK"

# ── build ONE frozen bot from HEAD ───────────────────────────────────────────
COMMIT="$(git -C "$REPO" rev-parse HEAD)"
FROZEN_DIR="$OUTDIR/frozen/ModularBot"
FROZEN_BIN="$FROZEN_DIR/ModularBot_bin"
FROZEN_JSON="$FROZEN_DIR/ModularBot.json"
mkdir -p "$FROZEN_DIR"
BUILDDIR="$WORK/head"
mkdir -p "$BUILDDIR"
echo "[ab] exporting HEAD ($COMMIT) -> $BUILDDIR"
git -C "$REPO" archive HEAD | tar -x -C "$BUILDDIR"

echo "[ab] building frozen ModularBot (nim c -d:release)…"
( cd "$BUILDDIR/ModularBot_garage" && \
  nim c -d:release --nimcache:"$OUTDIR/nimcache" --out:"$FROZEN_BIN" src/ModularBot.nim )
[[ -x "$FROZEN_BIN" ]] || { echo "ERROR: frozen build failed" >&2; exit 1; }
BINSHA="$(sha256sum "$FROZEN_BIN" | awk '{print $1}')"
echo "[ab] frozen binary sha256=$BINSHA"

# canonical bot identity for every arm (the dir is /named/ ModularBot)
cat > "$FROZEN_JSON" <<'JSON'
{
  "name": "ModularBot",
  "version": "0.1.0",
  "authors": ["Davide Cappellini"],
  "description": "frozen A/B build of ModularBot (ab_run.sh)",
  "homepage": "",
  "countryCodes": ["IT"],
  "gameTypes": ["classic", "1v1"],
  "platform": "Nim",
  "programmingLang": "Nim"
}
JSON

# ── session.json ─────────────────────────────────────────────────────────────
json_esc() { printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'; }
{
  printf '{\n'
  printf '  "commit": "%s",\n' "$COMMIT"
  printf '  "binary_sha256": "%s",\n' "$BINSHA"
  printf '  "binary": "frozen/ModularBot/ModularBot_bin",\n'
  printf '  "rounds": %d,\n' "$ROUNDS"
  printf '  "runs": %d,\n' "$RUNS"
  printf '  "conc": %d,\n' "$CONC"
  printf '  "timestamp": "%s",\n' "$(date -Is)"
  printf '  "outdir": "%s",\n' "$(json_esc "$OUTDIR")"
  printf '  "arms": [\n'
  for i in "${!ARM_NAMES[@]}"; do
    printf '    {"name": "%s", "env": "%s", "label": "%s"}' \
      "$(json_esc "${ARM_NAMES[$i]}")" "$(json_esc "${ARM_ENVS[$i]}")" "$(json_esc "${ARM_LABELS[$i]}")"
    (( i + 1 < ${#ARM_NAMES[@]} )) && printf ',' || true
    printf '\n'
  done
  printf '  ]\n}\n'
} > "$OUTDIR/session.json"

# ── one job (run in its own process group) ───────────────────────────────────
# exported so the detached `bash -c` subshell can call it; config via env
export REPO RUN_BRIDGE ROUNDS OUTDIR WORK FROZEN_BIN FROZEN_JSON

ab_run_one() {
  local arm="$1" run="$2" envspec="$3"
  local dir="$OUTDIR/$arm"
  local work="$WORK/$arm/run$run"
  local botdir="$work/bots/ModularBot"
  mkdir -p "$dir" "$botdir" "$work/drussgt" "$work/data"

  cp "$FROZEN_JSON" "$botdir/ModularBot.json"
  cat > "$botdir/ModularBot.sh" <<SH
#!/bin/sh
cd "\$(dirname "\$0")"
exec ./ModularBot_bin > "$dir/run$run.bot.stdout.log" 2> "$dir/run$run.bot.stderr.log"
SH
  chmod +x "$botdir/ModularBot.sh"
  ln -sf "$FROZEN_BIN" "$botdir/ModularBot_bin"

  rm -f "$dir/run$run.jsonl" "$dir/run$run.jsonl.rounds.json" \
        "$dir/run$run.jsonl.results.json" "$dir/run$run.events.jsonl" \
        "$dir/run$run.status"

  local rc=0
  # shellcheck disable=SC2086  # envspec is intentionally word-split
  env $envspec \
      DRUSSGT_BOTDIR="$work/drussgt/DrussGT" \
      DRUSSGT_DATA="$work/data" \
      AB_SESSION="$OUTDIR" \
      TR_EVENTS_OUT="$dir/run$run.events.jsonl" \
      timeout 900 "$RUN_BRIDGE" "$botdir" "$ROUNDS" "$dir/run$run.jsonl" \
      > "$dir/run$run.battle.log" 2>&1 || rc=$?
  echo "$rc" > "$dir/run$run.status"
  return 0
}
export -f ab_run_one

# ── cleanup: kill this session's children ────────────────────────────────────
JOB_PIDS=()
cleanup() {
  local rc=$?
  trap - EXIT INT TERM
  for p in "${JOB_PIDS[@]:-}"; do
    [[ -n "$p" ]] || continue
    kill -TERM -- "-$p" 2>/dev/null || kill -TERM "$p" 2>/dev/null || true
  done
  sleep 0.5
  for p in "${JOB_PIDS[@]:-}"; do
    [[ -n "$p" ]] || continue
    kill -KILL -- "-$p" 2>/dev/null || true
  done
  # backstop: any survivor whose cmdline references this session's outdir.
  # `[r]un_…` cannot match this script's own pkill command line.
  pkill -f "[r]un_bridge_battle.sh.*$OUTDIR" 2>/dev/null || true
  pkill -f "[r]obocode_shim.TrBattleCapture.*$OUTDIR" 2>/dev/null || true
  exit "$rc"
}
trap cleanup EXIT INT TERM

# ── launch the jobs, K at a time ─────────────────────────────────────────────
echo "[ab] session: ${#ARM_NAMES[@]} arms × $RUNS runs, rounds=$ROUNDS, conc=$CONC"
echo "[ab] outdir:  $OUTDIR"
RUNNING=0
for ai in "${!ARM_NAMES[@]}"; do
  arm="${ARM_NAMES[$ai]}"; envspec="${ARM_ENVS[$ai]}"
  for (( r=1; r<=RUNS; r++ )); do
    setsid bash -c 'ab_run_one "$1" "$2" "$3"' _ "$arm" "$r" "$envspec" &
    JOB_PIDS+=("$!")
    RUNNING=$((RUNNING + 1))
    if (( RUNNING >= CONC )); then
      wait -n || true
      RUNNING=$((RUNNING - 1))
    fi
  done
done
wait || true

# jobs are done — do not let the EXIT trap signal reused PIDs
JOB_PIDS=()

# ── aggregate status ─────────────────────────────────────────────────────────
FAILED=0
for ai in "${!ARM_NAMES[@]}"; do
  arm="${ARM_NAMES[$ai]}"
  for (( r=1; r<=RUNS; r++ )); do
    st="$(cat "$OUTDIR/$arm/run$r.status" 2>/dev/null || echo 999)"
    if [[ "$st" != "0" ]]; then
      echo "[ab] FAIL $arm run$r (rc=$st) — see $OUTDIR/$arm/run$r.battle.log" >&2
      FAILED=$((FAILED + 1))
    fi
  done
done

echo "[ab] done: $(( ${#ARM_NAMES[@]} * RUNS - FAILED )) ok, $FAILED failed"
echo "[ab] analyze with: python3 $REPO/tools/ab/ab_analyze.py $OUTDIR"
(( FAILED == 0 )) || exit 1