#!/usr/bin/env bash
# Pilot-gated, one-worker ECMWF System 51 soil-moisture baseline download.
set -u -o pipefail

project_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
start_month=${1:-1981-01}
end_month=${2:-2020-12}
retry_seconds=${ECMWF_SM_RETRY_SECONDS:-300}
retry_max_seconds=${ECMWF_SM_RETRY_MAX_SECONDS:-1800}
log_dir="$project_root/outputs/logs"
pilot_dir="$project_root/data/ecmwf/soil_moisture_pilot"
output_dir="$project_root/data/ecmwf/soil_moisture"
run_log="$log_dir/ecmwf_soil_moisture_download.log"
global_lock="$log_dir/ecmwf_baseline_download.lock"
pid_file="$log_dir/ecmwf_soil_moisture_download.pid"
pilot_failure_file="$log_dir/ecmwf_soil_moisture_pilot.last_failure"
mkdir -p "$log_dir" "$pilot_dir" "$output_dir"

exec 9>"$global_lock"
if ! flock -n 9; then
  printf '%s Soil-moisture worker cannot acquire global CDS lock; exiting.\n' "$(date -Is)" >>"$run_log"
  exit 75
fi
printf '%s\n' "$$" >"$pid_file"
trap 'rm -f "$pid_file"' EXIT

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

run_pilot() {
  local init=$1
  PYTHONPATH="$project_root/src" python3 -m dr_gcm.download_ecmwf_pilot \
    --init "$init" --group soil_moisture \
    --retry-partial --output-directory "$pilot_dir" >>"$run_log" 2>&1
}

printf '%s PILOT_START init=2005-01 leads=24,48,72\n' "$(date -Is)" >>"$run_log"
if ! run_pilot 2005-01; then
  date -Is >"$pilot_failure_file"
  exit 1
fi
printf '%s PILOT_SUCCESS init=2005-01\n' "$(date -Is)" >>"$run_log"
printf '%s PILOT_START init=1981-01 leads=24,48,72\n' "$(date -Is)" >>"$run_log"
if ! run_pilot 1981-01; then
  date -Is >"$pilot_failure_file"
  exit 1
fi
printf '%s PILOT_SUCCESS init=1981-01\n' "$(date -Is)" >>"$run_log"
rm -f "$pilot_failure_file"

current_month=$start_month
while [[ "$current_month" < "$end_month" || "$current_month" == "$end_month" ]]; do
  retry_delay=$retry_seconds
  while true; do
    printf '%s REQUEST init=%s group=soil_moisture\n' "$(date -Is)" "$current_month" >>"$run_log"
    if PYTHONPATH="$project_root/src" python3 -m dr_gcm.download_ecmwf_baseline \
      --retry-partial --init "$current_month" --group soil_moisture \
      --output-directory "$output_dir" >>"$run_log" 2>&1; then
      printf '%s SUCCESS init=%s group=soil_moisture\n' "$(date -Is)" "$current_month" >>"$run_log"
      break
    else
      status=$?
    fi
    printf '%s FAILED init=%s exit=%s retry_in=%ss\n' \
      "$(date -Is)" "$current_month" "$status" "$retry_delay" >>"$run_log"
    sleep "$retry_delay"
    retry_delay=$((retry_delay * 2))
    if (( retry_delay > retry_max_seconds )); then
      retry_delay=$retry_max_seconds
    fi
  done
  current_month=$(month_after "$current_month")
done
printf '%s COMPLETE range=%s..%s group=soil_moisture\n' \
  "$(date -Is)" "$start_month" "$end_month" >>"$run_log"