#!/bin/bash
set -euo pipefail

# Passive evaluator for saved fixed-sample WAVs.
# - Does NOT pause training (CPU-only ASR by default)
# - Logs metrics to W&B with eval/update as the x-axis
# - Appends raw JSONL to outputs/infer/fixed_sample_scores.jsonl

WORKSPACE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
RUNNER_CONTAINER="${RUNNER_CONTAINER:-cortex_runner_spark}"
PROJECT_NAME="${PROJECT_NAME:-cortazar_v2}"
SAMPLES_DIR="${SAMPLES_DIR:-/workspace/F5-TTS/ckpts/${PROJECT_NAME}/samples}"

CHECK_INTERVAL_SEC="${CHECK_INTERVAL_SEC:-120}"
EVAL_EVERY_UPDATES="${EVAL_EVERY_UPDATES:-1000}"
ASR_DEVICE="${ASR_DEVICE:-cpu}"
ASR_LANGUAGE="${ASR_LANGUAGE:-es}"
ASR_MODEL="${ASR_MODEL:-openai/whisper-small}"
WANDB_PREFIX="${WANDB_PREFIX:-val/fixed}"

STATE_FILE="$WORKSPACE_DIR/outputs/infer/eval_fixed_state.txt"
LOG_DIR="$WORKSPACE_DIR/outputs/logs"
LOG_FILE="$LOG_DIR/eval_fixed_daemon.log"
SCORES_FILE="$WORKSPACE_DIR/outputs/infer/fixed_sample_scores.jsonl"
PID_FILE="$WORKSPACE_DIR/outputs/infer/eval_fixed_daemon.pid"

mkdir -p "$WORKSPACE_DIR/outputs/infer"
mkdir -p "$LOG_DIR"
touch "$STATE_FILE"

if [[ -f "$PID_FILE" ]]; then
  old_pid="$(cat "$PID_FILE" 2>/dev/null || true)"
  if [[ -n "$old_pid" ]] && kill -0 "$old_pid" >/dev/null 2>&1; then
    echo "eval_fixed_samples_daemon already running (pid=$old_pid)" >&2
    exit 0
  fi
fi

echo $$ > "$PID_FILE"

cleanup() {
  rm -f "$PID_FILE" || true
}
trap cleanup EXIT INT TERM

log_msg() {
  echo "[$(date -u +%Y%m%dT%H%M%SZ)] $*" >> "$LOG_FILE"
}

LAST_EVAL=0
if [[ -s "$STATE_FILE" ]]; then
  read -r LAST_EVAL < "$STATE_FILE" || true
fi

next_candidates() {
  # List numeric updates for files update_<N>_fixed.wav
  docker exec "$RUNNER_CONTAINER" bash -lc "ls -1 ${SAMPLES_DIR}/update_*_fixed.wav 2>/dev/null | sed -E 's/.*update_([0-9]+)_fixed\\.wav/\\1/' | sort -n"
}

while true; do
  updates="$(next_candidates || true)"
  if [[ -n "$updates" ]]; then
    while read -r u; do
      [[ -z "$u" ]] && continue
      if (( u <= LAST_EVAL )); then
        continue
      fi
      if (( u % EVAL_EVERY_UPDATES != 0 )); then
        continue
      fi

      audio="${SAMPLES_DIR}/update_${u}_fixed.wav"
      ref_text="${SAMPLES_DIR}/update_${u}_fixed_gen_text.txt"
      if ! docker exec "$RUNNER_CONTAINER" bash -lc "test -f \"$audio\" && test -f \"$ref_text\""; then
        continue
      fi

      log_msg "Evaluating fixed sample (update=$u)"
      if docker exec "$RUNNER_CONTAINER" bash -lc \
        "python /workspace/scripts/pipeline/eval_audio_metrics.py --audio \"$audio\" --ref_text_file \"$ref_text\" --output_json \"/workspace${SCORES_FILE#$WORKSPACE_DIR}\" --step \"$u\" --wandb_project \"$PROJECT_NAME\" --wandb_prefix \"$WANDB_PREFIX\" --asr_model \"$ASR_MODEL\" --asr_device \"$ASR_DEVICE\" --language \"$ASR_LANGUAGE\" --skip_speaker" \
        >> "$LOG_FILE" 2>&1; then
        LAST_EVAL="$u"
        echo "$LAST_EVAL" > "$STATE_FILE"
      else
        log_msg "WARN: eval failed (update=$u)"
      fi
    done <<< "$updates"
  fi
  sleep "$CHECK_INTERVAL_SEC"
done