#!/bin/bash
# CFS prod 多起报时次 6h grib2 串行下载脚本
# Usage: ./cfs_dl.sh YYYYMMDD INIT_HOUR TOTAL_DAYS
# Example:
# 下载20260804 00时起报未来32天：./cfs_dl.sh 20260804 00 32
# 下载20260804 12时起报未来32天：./cfs_dl.sh 20260804 12 32

# ====================== 全局常量配置区 ======================
readonly PROGNAME="$(basename "$0")"
# NOAA官方prod数据源地址
readonly SRC_BASE_URL="https://nomads.ncep.noaa.gov/pub/data/nccf/com/cfs/prod"
readonly RETRY_TIMES=2           # wget单次网络重试次数
readonly CHECK_FILE_SIZE=100     # 极小文件判定阈值 100字节
readonly SAVE_ROOT="/data/public/Goldwind/cfsr/"

# 等待远端文件生成轮询控制（解决数据延迟产出）
readonly POLL_INTERVAL=120       # 每2分钟探测一次远端文件
readonly MAX_WAIT_SEC=$((40 * 60)) # 单文件最长等待40分钟，超时跳过该文件继续下一个

# 6小时预报时次
readonly TIME_STEPS=("00" "06" "12" "18")

# ====================== 底层下载函数（带超时防卡死） ======================
download() {
    local src_url="$1"
    local dst_file="$2"
    local filename="$3"

    echo "---------------------------------------------------------"
    echo "Start downloading: $filename"
    echo "URL: $src_url"

    wget --no-check-certificate -q -c \
        --dns-timeout=8 \
        --connect-timeout=10 \
        --read-timeout=15 \
        --timeout=20 \
        --tries="${RETRY_TIMES}" \
        -O "${dst_file}" "${src_url}"

    local ret=$?
    echo "Finish attempt: $filename wget exit code: ${ret}"
    return ${ret}
}

# ====================== 单文件轮询等待+校验下载核心 ======================
fetch_single_file() {
    local src_url="$1"
    local dst_file="$2"
    local filename="$3"

    mkdir -p "$(dirname "${dst_file}")"
    local start_ts=$(date +%s)
    local download_ok=0

    # 外层循环持续探测，直到下载成功或超时
    while true; do
        local now_ts=$(date +%s)
        local cost=$(( now_ts - start_ts ))

        # 超过最大等待时长，放弃当前文件，返回失败，继续下一个文件
        if [[ ${cost} -ge ${MAX_WAIT_SEC} ]]; then
            echo "ERROR: ${filename} wait over ${MAX_WAIT_SEC}s timeout, skip this file"
            return 1
        fi

        # 清理旧残缺文件
        rm -f "${dst_file}"
        # 执行单次下载
        download "${src_url}" "${dst_file}" "${filename}"
        local dl_ret=$?

        # 下载成功后校验文件大小
        if [[ ${dl_ret} -eq 0 && -f "${dst_file}" ]]; then
            local fsize
            fsize=$(stat -c %s "${dst_file}" 2>/dev/null || echo 0)
            if [[ ${fsize} -gt ${CHECK_FILE_SIZE} ]]; then
                echo "SUCCESS: ${filename} complete, size: ${fsize} bytes"
                download_ok=1
                break
            else
                echo "WARNING: ${filename} file too small(${fsize}B), wait and re-poll..."
            fi
        else
            echo "INFO: ${filename} remote file not ready or network error, sleep ${POLL_INTERVAL}s retry..."
        fi

        sleep ${POLL_INTERVAL}
    done

    [[ ${download_ok} -eq 1 ]] && return 0 || return 1
}

# ====================== 单起报时次批量下载 ======================
download_single_init_cycle() {
    local base_date="$1"
    local init_hour="$2"
    local total_fc_days="$3"

    local remote_subdir="cfs.${base_date}/${init_hour}/6hrly_grib_01"
    local full_prefix="${SRC_BASE_URL}/${remote_subdir}"
    local local_dir="${SAVE_ROOT}/${remote_subdir}"
    mkdir -p "${local_dir}"

    echo -e "\n===================== Init cycle: cfs.${base_date}/${init_hour} start download ====================="

    # 循环预报日：0 ~ total_fc_days-1
    for ((day_off=0; day_off < total_fc_days; day_off++)); do
        fc_ymd=$(date -d "${base_date} +${day_off} day" +%Y%m%d)
        echo -e "\n---------- Forecast date: ${fc_ymd} ----------"

        # 遍历4个6小时预报时次
        for step in "${TIME_STEPS[@]}"; do
            suffix="${fc_ymd}${step}.01.${base_date}${init_hour}.grb2"
            # pgbf 文件
            pgbf_name="pgbf${suffix}"
            pgbf_url="${full_prefix}/${pgbf_name}"
            pgbf_local="${local_dir}/${pgbf_name}"
            fetch_single_file "${pgbf_url}" "${pgbf_local}" "${pgbf_name}"

            # flxf 文件
            flxf_name="flxf${suffix}"
            flxf_url="${full_prefix}/${flxf_name}"
            flxf_local="${local_dir}/${flxf_name}"
            fetch_single_file "${flxf_url}" "${flxf_local}" "${flxf_name}"
        done
    done

    echo -e "\n===================== Init cycle cfs.${base_date}/${init_hour} all file process done ====================="
}

# ====================== 入口参数校验 & 主逻辑 ======================
main() {
    # 入参：脚本 日期 起报时次 预报总天数
    if [[ $# -ne 3 ]]; then
        echo "ERROR: Wrong arguments count!"
        echo "Usage: ${PROGNAME} YYYYMMDD INIT_HOUR TOTAL_DAYS"
        echo "Valid INIT_HOUR: 00 / 06 / 12 / 18"
        echo "Example 1: ${PROGNAME} 20260804 00 32"
        echo "Example 2: ${PROGNAME} 20260804 12 32"
        exit 1
    fi

    local ARG_DATE="$1"
    local ARG_HOUR="$2"
    local ARG_FC_DAYS="$3"

    # 1. 日期格式校验 8位数字
    if ! [[ "${ARG_DATE}" =~ ^[0-9]{8}$ ]]; then
        echo "ERROR: First param must be 8-digit date YYYYMMDD"
        exit 1
    fi

    # 2. 起报时次合法校验
    if ! [[ "${ARG_HOUR}" =~ ^(00|06|12|18)$ ]]; then
        echo "ERROR: Second param INIT_HOUR only support 00 / 06 / 12 / 18"
        exit 1
    fi

    # 3. 预报天数字符校验
    if ! [[ "${ARG_FC_DAYS}" =~ ^[0-9]+$ ]]; then
        echo "ERROR: Third param must be positive integer forecast days"
        exit 1
    fi
    local ARG_DATE_NEW=$(date +%Y%m%d)
    # 启动下载
    download_single_init_cycle "${ARG_DATE_NEW}" "${ARG_HOUR}" "${ARG_FC_DAYS}"
}

# 执行入口，透传命令行参数
main "$@"