#!/usr/bin/env bash
# Keep the approved one-worker ECMWF downloader and monthly builder alive.
# Usage: bash scripts/baseline_watchdog.sh [--once]
set -u -o pipefail

project_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
start_month=${BASELINE_START_MONTH:-1981-01}
end_month=${BASELINE_END_MONTH:-2020-12}
poll_seconds=${WATCHDOG_POLL_SECONDS:-300}
mode=${1:-}
log_dir="$project_root/outputs/logs"
watchdog_log="$log_dir/baseline_watchdog.log"
watchdog_lock="$log_dir/baseline_watchdog.lock"
watchdog_pid="$log_dir/baseline_watchdog.pid"
download_lock="$log_dir/ecmwf_baseline_download.lock"
monthly_lock="$log_dir/monthly_baseline.lock"

mkdir -p "$log_dir"
exec 8>"$watchdog_lock"
if ! flock -n 8; then
  printf '%s Another watchdog holds %s; exiting.\n' "$(date -Is)" "$watchdog_lock" >>"$watchdog_log"
  exit 75
fi
printf '%s\n' "$$" >"$watchdog_pid"
trap 'rm -f "$watchdog_pid"' EXIT

if [[ "$mode" != "" && "$mode" != "--once" ]]; then
  printf 'Usage: %s [--once]\n' "$0" >&2
  exit 2
fi
if ! [[ "$start_month" =~ ^[0-9]{4}-(0[1-9]|1[0-2])$ ]] ||
   ! [[ "$end_month" =~ ^[0-9]{4}-(0[1-9]|1[0-2])$ ]] ||
   ! [[ "$poll_seconds" =~ ^[1-9][0-9]*$ ]] ||
   [[ "$start_month" > "$end_month" ]]; then
  printf 'Invalid watchdog month range or WATCHDOG_POLL_SECONDS=%s\n' "$poll_seconds" >&2
  exit 2
fi

log() {
  printf '%s %s\n' "$(date -Is)" "$*" >>"$watchdog_log"
}

month_after() {
  date -d "${1}-01 +1 month" +%Y-%m
}

lock_is_free() {
  flock -n "$1" -c ':' >/dev/null 2>&1
}

first_missing_raw() {
  local current=$start_month instantaneous accumulated
  while [[ "$current" < "$end_month" || "$current" == "$end_month" ]]; do
    instantaneous="$project_root/data/ecmwf/raw_split/ecmwf_seas51_${current}_instantaneous.nc"
    accumulated="$project_root/data/ecmwf/raw_split/ecmwf_seas51_${current}_accumulated.nc"
    if [[ ! -f "$instantaneous" || ! -f "$accumulated" ]]; then
      printf '%s\n' "$current"
      return 0
    fi
    current=$(month_after "$current")
  done
  return 1
}

first_missing_monthly() {
  local current=$start_month output
  while [[ "$current" < "$end_month" || "$current" == "$end_month" ]]; do
    output="$project_root/data/ecmwf/monthly/ecmwf_seas51_${current}_monthly.nc"
    if [[ ! -f "$output" ]]; then
      printf '%s\n' "$current"
      return 0
    fi
    current=$(month_after "$current")
  done
  return 1
}

ensure_downloader() {
  local resume_month child
  if ! lock_is_free "$download_lock"; then
    return 0
  fi
  if ! resume_month=$(first_missing_raw); then
    log 'DOWNLOADER complete; no restart required'
    return 0
  fi
  log "DOWNLOADER absent; restarting range=${resume_month}..${end_month} workers=1"
  nohup setsid bash "$project_root/scripts/download_ecmwf_baseline.sh" \
    "$resume_month" "$end_month" 1 </dev/null \
    >>"$log_dir/ecmwf_baseline_download.log" 2>&1 &
  child=$!
  sleep 2
  if kill -0 "$child" 2>/dev/null || ! lock_is_free "$download_lock"; then
    log "DOWNLOADER restarted pid=${child} resume=${resume_month}"
  else
    log "DOWNLOADER restart failed pid=${child} resume=${resume_month}; will retry"
  fi
}

ensure_monthly() {
  local resume_month child
  if ! lock_is_free "$monthly_lock"; then
    return 0
  fi
  if ! resume_month=$(first_missing_monthly); then
    log 'MONTHLY complete; no restart required'
    return 0
  fi
  log "MONTHLY absent; restarting range=${resume_month}..${end_month}"
  nohup setsid bash "$project_root/scripts/build_monthly_baseline.sh" \
    "$resume_month" "$end_month" </dev/null \
    >>"$log_dir/monthly_baseline_supervisor.log" 2>&1 &
  child=$!
  sleep 2
  if kill -0 "$child" 2>/dev/null || ! lock_is_free "$monthly_lock"; then
    log "MONTHLY restarted pid=${child} resume=${resume_month}"
  else
    log "MONTHLY restart failed pid=${child} resume=${resume_month}; will retry"
  fi
}

log "WATCHDOG start pid=$$ range=${start_month}..${end_month} poll_seconds=${poll_seconds} mode=${mode:-loop}"
while true; do
  ensure_downloader
  ensure_monthly
  log 'HEARTBEAT checks_complete'
  if [[ "$mode" == "--once" ]]; then
    break
  fi
  sleep "$poll_seconds"
done