#!/usr/bin/env bash
set -euo pipefail

N_JOBS=62
CONFIG="axsmtg_config.example.yml"
SCRIPT="convert_axwliq_axwice_to_axsmtg.py"

LOG_ROOT="./logs_axsmtg_conversion"
RUN_LIST="${LOG_ROOT}/missing_xsmtg_axsmtg_dirs.txt"
SUCCESS_LIST="${LOG_ROOT}/success_dirs.txt"
FAILED_LIST="${LOG_ROOT}/failed_dirs.txt"

mkdir -p "$LOG_ROOT"
: > "$RUN_LIST"
: > "$SUCCESS_LIST"
: > "$FAILED_LIST"

echo "Scanning folders with T2MAX but without xsmtg/axsmtg..."

find . -type f -name "*T2MAX_daily.nc" | while read -r f; do
    dir=$(dirname "$f")

    # If xsmtg or axsmtg already exists, skip
    find "$dir" -maxdepth 1 -type f -name "*xsmtg_daily.nc"  | grep -q . && continue
    find "$dir" -maxdepth 1 -type f -name "*axsmtg_daily.nc" | grep -q . && continue

    echo "$dir"
done | sort -u > "$RUN_LIST"

N_TOTAL=$(wc -l < "$RUN_LIST")
echo "Found $N_TOTAL folders needing axsmtg conversion."
echo "Run list: $RUN_LIST"
echo "Logs:     $LOG_ROOT"
echo

if [[ "$N_TOTAL" -eq 0 ]]; then
    echo "Nothing to run."
    exit 0
fi

export SCRIPT CONFIG LOG_ROOT SUCCESS_LIST FAILED_LIST

run_one() {
    dir="$1"

    # Make a safe log filename from path:
    # ./2008/200805/20080516 -> 2008_200805_20080516
    safe_name=$(echo "$dir" | sed 's#^\./##; s#/#_#g')

    out_log="${LOG_ROOT}/${safe_name}.out.log"
    err_log="${LOG_ROOT}/${safe_name}.err.log"

    echo "[$(date '+%F %T')] START $dir" > "$out_log"

    if python "$SCRIPT" \
        --run-dir "$dir" \
        --config "$CONFIG" \
        --overwrite \
        >> "$out_log" 2> "$err_log"; then

        echo "[$(date '+%F %T')] SUCCESS $dir" >> "$out_log"
        echo "$dir" >> "$SUCCESS_LIST"
    else
        status=$?
        echo "[$(date '+%F %T')] FAILED $dir exit_code=$status" >> "$out_log"
        echo "$dir" >> "$FAILED_LIST"
        return "$status"
    fi
}

export -f run_one

echo "Running conversion with $N_JOBS parallel jobs..."
echo

cat "$RUN_LIST" | xargs -I {} -P "$N_JOBS" bash -c 'run_one "$@"' _ {}

echo
echo "Done."
echo "Success: $(wc -l < "$SUCCESS_LIST")"
echo "Failed:  $(wc -l < "$FAILED_LIST")"
echo
echo "Success list: $SUCCESS_LIST"
echo "Failed list:  $FAILED_LIST"

if [[ -s "$FAILED_LIST" ]]; then
    echo
    echo "Some jobs failed. Check logs under:"
    echo "  $LOG_ROOT"
fi