#!/usr/bin/env bash
# Lightweight heartbeat monitor for a long-running adaptive benchmark.
set -uo pipefail

if (( $# < 2 )); then
  echo "Usage: $0 RUN_ROOT MAIN_TMUX [ENGINE_PORT] [SERVICE_NAME_PATTERN]" >&2
  exit 2
fi

RUN_ROOT="$1"
MAIN_TMUX="$2"
ENGINE_PORT="${3:-}"
SERVICE_NAME_PATTERN="${4:-}"
INTERVAL_S="${HEARTBEAT_INTERVAL_S:-60}"
LOG_FILE="${HEARTBEAT_LOG_FILE:-${RUN_ROOT}/heartbeat.log}"

if ! [[ "$INTERVAL_S" =~ ^[1-9][0-9]*$ ]]; then
  echo "HEARTBEAT_INTERVAL_S must be a positive integer" >&2
  exit 2
fi

mkdir -p "$(dirname "$LOG_FILE")"

line_count() {
  local path="$1"
  if [[ -f "$path" ]]; then
    wc -l < "$path" | tr -d ' '
  else
    printf '0'
  fi
}

gpu_summary() {
  nvidia-smi \
    --query-gpu=index,memory.used,utilization.gpu \
    --format=csv,noheader,nounits 2>/dev/null | \
    awk -F',' '
      {
        for (i = 1; i <= NF; i++) gsub(/^[[:space:]]+|[[:space:]]+$/, "", $i)
        printf "%s%s:%sMiB/%s%%", (NR == 1 ? "" : ";"), $1, $2, $3
      }
      END { if (NR == 0) printf "unavailable" }
    '
}

heartbeat_once() {
  local timestamp tmux_state health_state container_names container_state
  local points shapes last_line gpu

  timestamp="$(date --iso-8601=seconds)"
  if tmux has-session -t "$MAIN_TMUX" 2>/dev/null; then
    tmux_state="up"
  else
    tmux_state="down"
  fi

  health_state="n/a"
  if [[ -n "$ENGINE_PORT" ]]; then
    if curl --fail --silent --max-time 5 \
      "http://127.0.0.1:${ENGINE_PORT}/health" >/dev/null 2>&1; then
      health_state="ok"
    else
      health_state="down"
    fi
  fi

  container_names=""
  if [[ -n "$SERVICE_NAME_PATTERN" ]]; then
    container_names="$(
      docker ps --filter "name=${SERVICE_NAME_PATTERN}" --format '{{.Names}}' 2>/dev/null | \
        paste -sd, -
    )"
  fi
  if [[ -n "$container_names" ]]; then
    container_state="$container_names"
  else
    container_state="none"
  fi

  points="$(line_count "${RUN_ROOT}/adaptive_points.jsonl")"
  shapes="$(line_count "${RUN_ROOT}/adaptive_shapes.jsonl")"
  last_line="$(
    tail -n 1 "${RUN_ROOT}/logs/orchestrator.log" 2>/dev/null | \
      tr '\t\r\n' ' ' | cut -c1-320
  )"
  gpu="$(gpu_summary)"

  printf '[%s] benchmark_tmux=%s health=%s service=%s points=%s shapes=%s gpu=%s last="%s"\n' \
    "$timestamp" "$tmux_state" "$health_state" "$container_state" \
    "$points" "$shapes" "$gpu" "$last_line" | tee -a "$LOG_FILE"

  [[ "$tmux_state" == "up" ]]
}

printf '[%s] heartbeat_started interval_s=%s main_tmux=%s run_root=%s\n' \
  "$(date --iso-8601=seconds)" "$INTERVAL_S" "$MAIN_TMUX" "$RUN_ROOT" | \
  tee -a "$LOG_FILE"

while heartbeat_once; do
  sleep "$INTERVAL_S"
done

printf '[%s] heartbeat_stopped reason=benchmark_tmux_down\n' \
  "$(date --iso-8601=seconds)" | tee -a "$LOG_FILE"