#!/usr/bin/env bash
# Download ECMWF SEAS5.1 baseline initial months concurrently and resumably.
# Usage: bash scripts/download_ecmwf_baseline.sh [START_YYYY-MM] [END_YYYY-MM] [WORKERS]
# Each worker gets a disjoint sequence of months; completed files are skipped.
# A rejected CDS request is retried for the *same* month with backoff. A month
# that repeatedly reports an ECMWF damaged tape is explicitly deferred in a
# provenance ledger so later months can continue; it is never silently skipped.
# One worker is the safe default for this CDS account; pass an explicit worker
# count only after confirming that CDS accepts it.
set -u -o pipefail

project_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
start_month=${1:-1981-01}
end_month=${2:-2020-12}
workers=${3:-1}
retry_seconds=${ECMWF_RETRY_SECONDS:-300}
retry_max_seconds=${ECMWF_RETRY_MAX_SECONDS:-1800}
damaged_tape_retry_limit=${ECMWF_DAMAGED_TAPE_RETRY_LIMIT:-3}
log_dir="$project_root/outputs/logs"
mkdir -p "$log_dir"
run_log="$log_dir/ecmwf_baseline_download.log"
deferred_gap_file="$log_dir/ecmwf_baseline_deferred_gaps.tsv"
lock_file="$log_dir/ecmwf_baseline_download.lock"
pid_file="$log_dir/ecmwf_baseline_download.pid"
exec 9>"$lock_file"
if ! flock -n 9; then
  printf '%s Another baseline worker holds %s; exiting.\n' "$(date -Is)" "$lock_file" >>"$run_log"
  exit 75
fi
printf '%s\n' "$$" >"$pid_file"
trap 'rm -f "$pid_file"' EXIT

if ! [[ "$workers" =~ ^[1-9][0-9]*$ ]] || ! [[ "$retry_seconds" =~ ^[1-9][0-9]*$ ]] || ! [[ "$retry_max_seconds" =~ ^[1-9][0-9]*$ ]] || ! [[ "$damaged_tape_retry_limit" =~ ^[1-9][0-9]*$ ]]; then
  printf 'WORKERS and retry settings must be positive integers; got workers=%s retry=%s retry_max=%s damaged_tape_limit=%s\n' \
    "$workers" "$retry_seconds" "$retry_max_seconds" "$damaged_tape_retry_limit" >&2
  exit 2
fi

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

download_worker() {
  local worker=$1 current_month status retry_delay manifest_path
  local damaged_tape_retries deferred
  current_month=$(month_after "$start_month" "$worker")
  while [[ "$current_month" < "$end_month" || "$current_month" == "$end_month" ]]; do
    printf '%s WORKER=%s REQUEST init=%s\n' "$(date -Is)" "$worker" "$current_month" >>"$run_log"
    retry_delay=$retry_seconds
    damaged_tape_retries=0
    deferred=0
    manifest_path="$project_root/data/ecmwf/raw_split/ecmwf_seas51_${current_month}_manifest.json"
    while true; do
      if PYTHONPATH="$project_root/src" python3 -m dr_gcm.download_ecmwf_baseline --retry-partial --init "$current_month" >>"$run_log" 2>&1; then
        break
      else
        status=$?
        if [[ -f "$manifest_path" ]] && grep -qi 'damaged tape' "$manifest_path"; then
          damaged_tape_retries=$((damaged_tape_retries + 1))
          printf '%s WORKER=%s DAMAGED_TAPE init=%s attempt=%s/%s\n' \
            "$(date -Is)" "$worker" "$current_month" "$damaged_tape_retries" "$damaged_tape_retry_limit" >>"$run_log"
          if (( damaged_tape_retries >= damaged_tape_retry_limit )); then
            if [[ ! -e "$deferred_gap_file" ]]; then
              printf 'timestamp\tinitial_month\treason\tattempts\n' >>"$deferred_gap_file"
            fi
            printf '%s\t%s\tdamaged_tape\t%s\n' \
              "$(date -Is)" "$current_month" "$damaged_tape_retries" >>"$deferred_gap_file"
            printf '%s WORKER=%s DEFERRED init=%s reason=damaged_tape attempts=%s ledger=%s\n' \
              "$(date -Is)" "$worker" "$current_month" "$damaged_tape_retries" "$deferred_gap_file" >>"$run_log"
            deferred=1
            break
          fi
        else
          damaged_tape_retries=0
        fi
        printf '%s WORKER=%s FAILED init=%s exit=%s; retrying in %ss\n' "$(date -Is)" "$worker" "$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
        printf '%s WORKER=%s RETRY init=%s\n' "$(date -Is)" "$worker" "$current_month" >>"$run_log"
      fi
    done
    if (( deferred == 1 )); then
      current_month=$(month_after "$current_month" "$workers")
      continue
    fi
    printf '%s WORKER=%s SUCCESS init=%s\n' "$(date -Is)" "$worker" "$current_month" >>"$run_log"
    current_month=$(month_after "$current_month" "$workers")
  done
}

printf '%s START range=%s..%s workers=%s\n' "$(date -Is)" "$start_month" "$end_month" "$workers" >>"$run_log"
for ((worker = 0; worker < workers; worker++)); do
  download_worker "$worker" &
done
wait
printf '%s COMPLETE range=%s..%s workers=%s\n' "$(date -Is)" "$start_month" "$end_month" "$workers" >>"$run_log"