#!/usr/bin/env bash

set -Eeuo pipefail

# Input root used by the legacy CLWRF and RegCM workflow. Downloaded data for a
# cycle must be in $OUTPUT_DIR/YYYYMMDD.
OUTPUT_DIR="${OUTPUT_DIR:-/home/nampq/data/CFS_Ope/00}"
LEGACY_CLWRF_DIR="${LEGACY_CLWRF_DIR:-/work/users/nampq/working/Proj_T.Tan/Operational/Seasonal_FC/CLWRF/CFS}"
LEGACY_REGCM_DIR="${LEGACY_REGCM_DIR:-/work/users/nampq/working/Proj_T.Tan/Operational/Seasonal_FC/RegCM/CFS}"
LEGACY_DATA_DIR="${LEGACY_DATA_DIR:-/home/nampq/data/CFS_Ope/00}"
CURRENT_RESULT_ROOT="${CURRENT_RESULT_ROOT:-/work/users/nampq/working/Proj_T.Tan/Operational/Seasonal_FC/data/current}"
execute=0
rebuild=0

usage() {
    cat <<'EOF'
Usage:
  ./run_rcms_manual.sh [--execute] [--rebuild-preprocessing] YEAR MONTH DAY [HOUR]

The default mode validates the CFSv2 input and prints the old CLWRF + RegCM
workflow without running it. Add --execute only after reviewing the plan.

  --rebuild-preprocessing
      With --execute, move existing state for this cycle to timestamped backup
      directories and rebuild WPS, RegCM Dataset, and IC/BC from raw CFSv2 GRIB2.
EOF
}

die() {
    printf 'ERROR: %s\n' "$*" >&2
    exit 1
}

while (($#)); do
    case "$1" in
        --execute) execute=1; shift ;;
        --rebuild-preprocessing) rebuild=1; shift ;;
        -h|--help) usage; exit 0 ;;
        --) shift; break ;;
        -*) die "unknown option: $1" ;;
        *) break ;;
    esac
done

(($# == 3 || $# == 4)) || { usage >&2; exit 2; }
year=$1
month=$2
day=$3
hour=${4:-00}

[[ $year =~ ^[0-9]{4}$ ]] || die "YEAR must contain four digits"
[[ $month =~ ^[0-9]{1,2}$ ]] || die "MONTH must contain one or two digits"
[[ $day =~ ^[0-9]{1,2}$ ]] || die "DAY must contain one or two digits"
[[ $hour =~ ^[0-9]{1,2}$ ]] || die "HOUR must contain one or two digits"
printf -v month '%02d' "$((10#$month))"
printf -v day '%02d' "$((10#$day))"
printf -v hour '%02d' "$((10#$hour))"
ymd="${year}${month}${day}"
ymdh="${ymd}${hour}"

normalized=$(date -u -d "${year}-${month}-${day} ${hour}:00:00" +%Y%m%d%H 2>/dev/null) ||
    die "invalid UTC date: ${year}-${month}-${day} ${hour}:00"
[[ $normalized == "$ymdh" ]] || die "invalid UTC date: ${year}-${month}-${day} ${hour}:00"

OUTPUT_DIR=$(realpath -m -- "$OUTPUT_DIR")
LEGACY_DATA_DIR=$(realpath -m -- "$LEGACY_DATA_DIR")
cycle_dir="${OUTPUT_DIR}/${ymd}"
[[ -d $cycle_dir ]] || die "download directory does not exist: $cycle_dir"
for required_script in \
    "$LEGACY_CLWRF_DIR/run_wps.sh" \
    "$LEGACY_CLWRF_DIR/run_wrf.sh" \
    "$LEGACY_CLWRF_DIR/run_post-process.sh" \
    "$LEGACY_REGCM_DIR/run_icbc.sh" \
    "$LEGACY_REGCM_DIR/run_regcm.sh" \
    "$LEGACY_REGCM_DIR/run_post-process.sh"; do
    [[ -x $required_script ]] || die "missing executable legacy component: $required_script"
done
command -v sbatch >/dev/null 2>&1 || die "sbatch is not available"

shopt -s nullglob
flux_files=("$cycle_dir"/flxf*.grb2)
pressure_files=("$cycle_dir"/pgbf*.grb2)
shopt -u nullglob
((${#flux_files[@]} > 0)) || die "no flxf GRIB2 files in $cycle_dir"
((${#pressure_files[@]} > 0)) || die "no pgbf GRIB2 files in $cycle_dir"

declare -A pressure_names=()
declare -A flux_times=()
for path in "${pressure_files[@]}"; do
    [[ -s $path ]] || die "empty pressure file: $path"
    pressure_names["${path##*/}"]=1
done
for path in "${flux_files[@]}"; do
    [[ -s $path ]] || die "empty flux file: $path"
    filename=${path##*/}
    [[ $filename =~ ^flxf([0-9]{10})\.01\.([0-9]{10})\.grb2$ ]] ||
        die "unexpected CFS filename: $filename"
    [[ ${BASH_REMATCH[2]} == "$ymdh" ]] ||
        die "file initialization ${BASH_REMATCH[2]} does not match requested cycle $ymdh"
    flux_times["${BASH_REMATCH[1]}"]=1
    counterpart="pgbf${filename#flxf}"
    [[ ${pressure_names[$counterpart]+yes} ]] || die "missing pair for $filename"
done
((${#flux_files[@]} == ${#pressure_files[@]})) ||
    die "flxf/pgbf counts differ: ${#flux_files[@]} versus ${#pressure_files[@]}"

IFS=$'\n' sorted_flux=($(printf '%s\n' "${flux_files[@]}" | sort))
unset IFS
first_name=${sorted_flux[0]##*/}
last_name=${sorted_flux[${#sorted_flux[@]}-1]##*/}
[[ $first_name =~ ^flxf([0-9]{10})\.01\.([0-9]{10})\.grb2$ ]] ||
    die "unexpected CFS filename: $first_name"
start_date=${BASH_REMATCH[1]}
initial_date=${BASH_REMATCH[2]}
[[ $initial_date == "$ymdh" ]] ||
    die "file initialization $initial_date does not match requested cycle $ymdh"
[[ $last_name =~ ^flxf([0-9]{10})\.01\.([0-9]{10})\.grb2$ ]] ||
    die "unexpected CFS filename: $last_name"
end_date=${BASH_REMATCH[1]}

expected_time=$start_date
expected_count=0
while [[ $expected_time -le $end_date ]]; do
    [[ ${flux_times[$expected_time]+yes} ]] ||
        die "missing six-hour CFS pair at $expected_time"
    expected_time=$(date -u -d \
        "${expected_time:0:4}-${expected_time:4:2}-${expected_time:6:2} ${expected_time:8:2}:00 UTC 6 hours" \
        +%Y%m%d%H)
    ((expected_count += 1))
done
((expected_count == ${#flux_files[@]})) ||
    die "unexpected timestamps outside the continuous $start_date-$end_date interval"

result_date="${year}-${month}-${day}"
clwrf_stage1=0
clwrf_stage2=0
regcm_stage1=0
regcm_stage2=0

# GD2 markers follow the legacy run_check.sh scripts. Also accept NetCDF files
# still present in done/<S_DATE>, before they are copied/moved to data/current.
if [[ -s "${CURRENT_RESULT_ROOT}/CLWRF/${start_date}/rainc.nc" ]] ||
    [[ -n $(find "${LEGACY_CLWRF_DIR}/done/${start_date}" -type f -name '*.nc' -print -quit 2>/dev/null) ]]; then
    clwrf_stage2=1
    clwrf_stage1=1
fi
if [[ -s "${CURRENT_RESULT_ROOT}/RegCM/${start_date}/daily/pr.nc" ]] ||
    [[ -n $(find "${LEGACY_REGCM_DIR}/done/${start_date}" -type f -name '*.nc' -print -quit 2>/dev/null) ]]; then
    regcm_stage2=1
    regcm_stage1=1
fi

# GD1 means the numerical integration reached its successful end. CLWRF clean
# preserves wrfout files in wdir2; RegCM records an explicit success line and
# moves the run tree to wdir2 before post-processing.
if [[ -n $(find "${LEGACY_CLWRF_DIR}/wdir2/${start_date}/output" -type f -name 'wrfout_*' -print -quit 2>/dev/null) ]] ||
    grep -q 'SUCCESS COMPLETE WRF' "${LEGACY_CLWRF_DIR}/wdir/${start_date}/rsl.error.0000" 2>/dev/null; then
    clwrf_stage1=1
fi
if [[ -d "${LEGACY_REGCM_DIR}/wdir2/${start_date}" ]] ||
    grep -q 'RegCM V4 simulation successfully reached end' "${LEGACY_REGCM_DIR}/logs_run_${start_date}" 2>/dev/null; then
    regcm_stage1=1
fi

if ((clwrf_stage2)); then
    printf 'CLWRF has been run for %s\n' "$result_date"
elif ((clwrf_stage1)); then
    printf 'CLWRF GD1 (integration) has completed for %s; GD2 (post-processing) is pending.\n' "$result_date"
else
    printf 'CLWRF GD1 (integration) is pending for %s.\n' "$result_date"
fi

if ((regcm_stage2)); then
    printf 'RegCM has been run for %s\n' "$result_date"
elif ((regcm_stage1)); then
    printf 'RegCM GD1 (integration) has completed for %s; GD2 (post-processing) is pending.\n' "$result_date"
else
    printf 'RegCM GD1 (integration) is pending for %s.\n' "$result_date"
fi

if ((clwrf_stage2 && regcm_stage2)); then
    printf 'All RCM results are already available; no model command was run.\n'
    exit 0
fi

printf '%s\n' \
    "RCM manual workflow" \
    "  requested cycle : $ymdh UTC" \
    "  CFS input       : $cycle_dir" \
    "  paired files    : ${#flux_files[@]} flxf + ${#pressure_files[@]} pgbf" \
    "  model interval  : $start_date to $end_date" \
    "  CLWRF work dir  : $LEGACY_CLWRF_DIR" \
    "  RegCM work dir  : $LEGACY_REGCM_DIR"

[[ $OUTPUT_DIR == "$LEGACY_DATA_DIR" ]] ||
    die "legacy drivers read $LEGACY_DATA_DIR, but OUTPUT_DIR is $OUTPUT_DIR"

printf '%s\n' \
    "Workflow:" \
    "  1. Handle CLWRF for $start_date only: WPS -> sbatch run_wrf.sh." \
    "  2. Handle RegCM for $start_date only: GRIB2 -> IC/BC -> sbatch run_regcm.sh." \
    "  3. Resume post-processing when GD1 is complete but GD2 is missing." \
    "  4. Never call either legacy autorun.sh."

if ((execute == 0)); then
    if ((rebuild)); then
        printf 'PLAN ONLY: preprocessing state for %s would be backed up and rebuilt.\n' "$start_date"
    fi
    printf 'PLAN ONLY: no model command was run. Add --execute to launch the legacy workflow.\n'
    exit 0
fi

if ((rebuild)); then
    if command -v squeue >/dev/null 2>&1 &&
        squeue -u "$(id -un)" -h -n clWRF_run,RegCM_run | grep -q .; then
        die "an CLWRF or RegCM Slurm job is active; refusing to rebuild preprocessing"
    fi

    backup_tag=$(date -u +%Y%m%dT%H%M%SZ)
    clwrf_backup="${LEGACY_CLWRF_DIR}/manual_backup/${start_date}.${backup_tag}"
    regcm_backup="${LEGACY_REGCM_DIR}/manual_backup/${start_date}.${backup_tag}"
    clwrf_backup_created=0
    regcm_backup_created=0

    for backup_entry in \
        "wdir|${LEGACY_CLWRF_DIR}/wdir/${start_date}" \
        "wdir2|${LEGACY_CLWRF_DIR}/wdir2/${start_date}" \
        "done|${LEGACY_CLWRF_DIR}/done/${start_date}" \
        "output|${LEGACY_CLWRF_DIR}/output/${start_date}"; do
        backup_group=${backup_entry%%|*}
        path=${backup_entry#*|}
        if [[ -e $path ]]; then
            if ((clwrf_backup_created == 0)); then
                mkdir -p -- "$clwrf_backup"
                clwrf_backup_created=1
            fi
            printf 'Backing up CLWRF state: %s\n' "$path"
            mkdir -p -- "$clwrf_backup/$backup_group"
            mv -- "$path" "$clwrf_backup/$backup_group/"
        fi
    done

    for backup_entry in \
        "Dataset_CFS|${LEGACY_REGCM_DIR}/Dataset/CFS/${start_date}" \
        "wdir|${LEGACY_REGCM_DIR}/wdir/${start_date}" \
        "wdir2|${LEGACY_REGCM_DIR}/wdir2/${start_date}" \
        "done|${LEGACY_REGCM_DIR}/done/${start_date}" \
        "namelist|${LEGACY_REGCM_DIR}/namelist.in_${start_date}" \
        "logs|${LEGACY_REGCM_DIR}/logs_run_${start_date}"; do
        backup_group=${backup_entry%%|*}
        path=${backup_entry#*|}
        if [[ -e $path ]]; then
            if ((regcm_backup_created == 0)); then
                mkdir -p -- "$regcm_backup"
                regcm_backup_created=1
            fi
            printf 'Backing up RegCM state: %s\n' "$path"
            mkdir -p -- "$regcm_backup/$backup_group"
            mv -- "$path" "$regcm_backup/$backup_group/"
        fi
    done

    ((clwrf_backup_created == 0)) || printf 'CLWRF backup: %s\n' "$clwrf_backup"
    ((regcm_backup_created == 0)) || printf 'RegCM backup: %s\n' "$regcm_backup"
    clwrf_stage1=0
    clwrf_stage2=0
    regcm_stage1=0
    regcm_stage2=0
fi

if ((clwrf_stage2)); then
    printf 'CLWRF: GD2 already complete; nothing to run.\n'
elif ((clwrf_stage1)); then
    printf 'CLWRF: resuming GD2 post-processing.\n'
    (
        cd -- "$LEGACY_CLWRF_DIR"
        ./run_post-process.sh
    )
elif [[ -d "${LEGACY_CLWRF_DIR}/wdir/${start_date}" ]]; then
    printf 'CLWRF: wdir/%s exists; job may be queued, running, or failed. Not resubmitting.\n' "$start_date"
else
    printf 'CLWRF: preparing WPS and submitting integration for %s only.\n' "$start_date"
    (
        cd -- "$LEGACY_CLWRF_DIR"
        ./run_wps.sh "$start_date" "$end_date"
        sbatch run_wrf.sh "$start_date" "$end_date"
    )
fi

if ((regcm_stage2)); then
    printf 'RegCM: GD2 already complete; nothing to run.\n'
elif ((regcm_stage1)); then
    printf 'RegCM: resuming GD2 post-processing.\n'
    (
        cd -- "$LEGACY_REGCM_DIR"
        ./run_post-process.sh
    )
elif [[ -d "${LEGACY_REGCM_DIR}/wdir/${start_date}" ]]; then
    printf 'RegCM: wdir/%s exists; job may be queued, running, or failed. Not resubmitting.\n' "$start_date"
else
    regcm_dataset="${LEGACY_REGCM_DIR}/Dataset/CFS/${start_date}/01"
    sst_file="${regcm_dataset}/SST/sst.${start_date}.nc"
    required_regcm_files=(
        "$sst_file"
        "${regcm_dataset}/PLEV/air.${start_date}.nc"
        "${regcm_dataset}/PLEV/rhum.${start_date}.nc"
        "${regcm_dataset}/PLEV/uwnd.${start_date}.nc"
        "${regcm_dataset}/PLEV/vwnd.${start_date}.nc"
        "${regcm_dataset}/PLEV/hgt.${start_date}.nc"
    )
    dataset_ready=1
    for path in "${required_regcm_files[@]}"; do
        [[ -s $path ]] || dataset_ready=0
    done

    if [[ -d "${LEGACY_REGCM_DIR}/Dataset/CFS/${start_date}" ]] && ((dataset_ready == 0)); then
        die "partial RegCM Dataset exists for $start_date; review it before retrying"
    fi

    if ((dataset_ready == 0)); then
        printf 'RegCM: converting GRIB2 input for %s only.\n' "$start_date"
        (
            cd -- "$LEGACY_REGCM_DIR"
            # shellcheck disable=SC1091
            source ../../../configure_regcm
            command -v wgrib2 >/dev/null 2>&1 || die "wgrib2 is unavailable after configure_regcm"
            mkdir -p -- "$regcm_dataset/SST" "$regcm_dataset/PLEV"
            cmdline=(-nc_table "$LEGACY_REGCM_DIR/Dataset/prc.table" -append -nc4 -netcdf)
            for path in "${flux_files[@]}"; do
                wgrib2 "$path" -match 'TMP:surface' "${cmdline[@]}" "$sst_file"
            done
            for path in "${pressure_files[@]}"; do
                wgrib2 "$path" -match ':TMP'  -match 'mb:' "${cmdline[@]}" "${regcm_dataset}/PLEV/air.${start_date}.nc"
                wgrib2 "$path" -match ':RH'   -match 'mb:' "${cmdline[@]}" "${regcm_dataset}/PLEV/rhum.${start_date}.nc"
                wgrib2 "$path" -match ':UGRD' -match 'mb:' "${cmdline[@]}" "${regcm_dataset}/PLEV/uwnd.${start_date}.nc"
                wgrib2 "$path" -match ':VGRD' -match 'mb:' "${cmdline[@]}" "${regcm_dataset}/PLEV/vwnd.${start_date}.nc"
                wgrib2 "$path" -match ':HGT'  -match 'mb:' "${cmdline[@]}" "${regcm_dataset}/PLEV/hgt.${start_date}.nc"
            done
        )
    else
        printf 'RegCM: converted Dataset already exists for %s.\n' "$start_date"
    fi

    printf 'RegCM: preparing IC/BC and submitting integration for %s only.\n' "$start_date"
    (
        cd -- "$LEGACY_REGCM_DIR"
        export S_DATE=$start_date
        export E_DATE=$end_date
        ./run_icbc.sh "$start_date" "$end_date"
        sbatch --export=ALL,S_DATE="$start_date",E_DATE="$end_date" \
            run_regcm.sh "$start_date" "$end_date"
    )
fi