#!/bin/bash
#set -x

readonly INPUT_DATE="$1"      # 格式：20250701
readonly INPUT_HOUR="$2"      # 格式：12
readonly INPUT_END_COUNT="$3" # 格式：168
readonly PROGNAME="$(basename $0)"
readonly SRC_BASE_URL="https://osdf-director.osg-htc.org/ncar/gdex/d084001"
readonly SAVE_PATH="/data/dfs"  # 修改为实际的保存路径
readonly PREFIX="gfs.0p25."
readonly SUFFIX=".grib2"
readonly RETRYS=10
readonly CHECK_FILE_SIZE="1000000"  # 根据实际文件大小调整

# 提取年月日
extract_date_parts() {
    local date_str="$1"
    local year="${date_str:0:4}"
    local month="${date_str:4:2}"
    local day="${date_str:6:2}"
    echo "${year} ${month} ${day}"
}

download(){
    local src_path="$1"
    local dst_path="$2"
    local dst_file="$3"
    
    echo "Downloading $src_path ..."
    wget --no-check-certificate -q -O "$dst_path/$dst_file" "$src_path"
    local wget_status=$?
    
    if [ $wget_status -ne 0 ]; then
        echo "Warning: wget failed with status $wget_status for $src_path"
        return 1
    fi
    return 0
}

download_gfs_grib(){
    local date_str="$1"
    local hour="$2"
    local tim="$3"
    
    # 解析日期
    read year month day <<< $(extract_date_parts $date_str)
    
    # 构造文件名和URL
    local src_file_name="${PREFIX}${date_str}${hour}.f$(printf "%03d" $tim)${SUFFIX}"
    local src_file_url="${SRC_BASE_URL}/${year}/${date_str}/${src_file_name}"
    
    # 目标存储路径（根据你的实际结构）
    local dst_dir="${SAVE_PATH}/dfs_${date_str}/${hour}/atmos"
    local dst_file="${dst_dir}/${src_file_name}"
    
    # 创建目标目录
    mkdir -p "$dst_dir"
    
    # 检查文件是否已存在且完整
    if [[ -f "$dst_file" ]]; then
        fs=$(stat -c %s "$dst_file" 2>/dev/null || echo 0)
        if [[ $fs -ge $CHECK_FILE_SIZE ]]; then
            echo "File $src_file_name already exists and is complete, skipping."
            return 0
        else
            echo "File $src_file_name exists but size is $fs (min expected: $CHECK_FILE_SIZE), re-downloading."
            rm -f "$dst_file"
        fi
    fi
    
    # 下载文件
    download "$src_file_url" "$dst_dir" "$src_file_name"
    
    # 检查文件大小，重试机制
    for i in $(seq $RETRYS)
    do 
        if [[ -f "$dst_file" ]]; then
            fs=$(stat -c %s "$dst_file" 2>/dev/null || echo 0)
            if [[ $fs -lt $CHECK_FILE_SIZE ]]; then
                sleep 30
                echo "Retry $i/$RETRYS for $src_file_name (size: $fs)"
                download "$src_file_url" "$dst_dir" "$src_file_name"
            else
                echo "Download completed: $src_file_name (size: $fs bytes)"
                return 0
            fi
        else
            sleep 30
            echo "Retry $i/$RETRYS for $src_file_name (file not found)"
            download "$src_file_url" "$dst_dir" "$src_file_name"
        fi
    done
    
    echo "Error: Failed to download $src_file_name after $RETRYS retries"
    return 1
}

download_gfs_files(){
    local date_str="$1"
    local hour="$2"
    local end_count="$3"
    
    echo "Starting download for date: $date_str, hour: $hour, forecast hours: 0-$end_count"
    
    local failed_files=0
    local success_files=0
    
    # 生成从0到end_count的序列，步长为3（0,3,6,...,168）
    local max_count=$((end_count / 3 * 3))
    for ((t=0; t<=max_count; t+=3)); do
        echo "Downloading forecast hour $(printf "%03d" $t) of $max_count"
        download_gfs_grib "$date_str" "$hour" "$t"
        if [ $? -ne 0 ]; then
            let failed_files++
            echo "Failed to download forecast hour $(printf "%03d" $t)"
        else
            let success_files++
        fi
        sleep 1  # 短暂延迟，避免请求过快
    done
    
    echo "Download summary: $success_files successful, $failed_files failed"
    return $failed_files
}

check_download_status(){
    local date_str="$1"
    local hour="$2"
    local end_count="$3"
    
    local dst_dir="${SAVE_PATH}/dfs_${date_str}/${hour}/atmos"
    local status_code=0
    local expected_count=$((end_count / 3 + 1))  # 包含f000
    
    echo "Checking download status for $date_str $hour..."
    echo "Expected files: $expected_count (f000 to f$(printf "%03d" $end_count) step 3)"
    
    local downloaded_count=0
    
    for ((t=0; t<=end_count; t+=3)); do
        local tim=$(printf "%03d" $t)
        local file_name="${PREFIX}${date_str}${hour}.f${tim}${SUFFIX}"
        local dst_file="${dst_dir}/${file_name}"
        
        # 检查文件是否存在
        if [[ ! -f "$dst_file" ]]; then
            echo "Missing file: $file_name"
            let status_code++
            continue
        fi
        
        # 检查文件大小
        fs=$(stat -c %s "$dst_file" 2>/dev/null || echo 0)
        if [[ $fs -lt $CHECK_FILE_SIZE ]]; then
            echo "Incomplete file: $file_name (size: $fs bytes)"
            let status_code++
        else
            let downloaded_count++
            echo "OK: $file_name (size: $fs bytes)"
        fi
    done
    
    echo "Downloaded $downloaded_count/$expected_count files"
    echo "Total missing/incomplete files: $status_code"
    return $status_code
}

generate_date_range() {
    local start_date="$1"
    local end_date="$2"
    local hour="$3"
    local end_count="$4"
    
    # 转换为时间戳
    local start_ts=$(date -d "$start_date" +%s 2>/dev/null || date -j -f "%Y%m%d" "$start_date" +%s 2>/dev/null)
    local end_ts=$(date -d "$end_date" +%s 2>/dev/null || date -j -f "%Y%m%d" "$end_date" +%s 2>/dev/null)
    
    if [ -z "$start_ts" ] || [ -z "$end_ts" ]; then
        echo "Error: Invalid date format. Please use YYYYMMDD format."
        exit 1
    fi
    
    local current_ts=$start_ts
    local success_count=0
    local fail_count=0
    
    while [ $current_ts -le $end_ts ]; do
        local current_date=$(date -d "@$current_ts" +%Y%m%d 2>/dev/null || date -j -f "%s" "$current_ts" +%Y%m%d 2>/dev/null)
        
        echo "========================================"
        echo "Processing date: $current_date"
        echo "========================================"
        
        # 下载文件
        download_gfs_files "$current_date" "$hour" "$end_count"
        local download_result=$?
        
        # 检查下载状态
        check_download_status "$current_date" "$hour" "$end_count"
        local check_result=$?
        
        if [ $download_result -eq 0 ] && [ $check_result -eq 0 ]; then
            echo "Successfully downloaded all files for $current_date"
            let success_count++
        else
            echo "Warning: Some files failed for $current_date"
            let fail_count++
        fi
        
        # 增加一天
        current_ts=$((current_ts + 86400))
        sleep 2  # 日期间延迟
    done
    
    echo "========================================"
    echo "Summary:"
    echo "  Successfully processed dates: $success_count"
    echo "  Dates with failures: $fail_count"
    echo "========================================"
}

main(){
    # 用法1: 下载单日数据
    # ./script.sh 20250701 12 168
    if [[ $# -eq 3 ]]; then
        local date_str="$1"
        local hour="$2"
        local end_count="$3"
        
        # 验证日期格式
        if ! [[ $date_str =~ ^[0-9]{8}$ ]]; then
            echo "Error: Date format should be YYYYMMDD"
            exit 1
        fi
        
        # 验证小时
        if ! [[ $hour =~ ^(00|06|12|18)$ ]]; then
            echo "Warning: Hour $hour may not be valid (typical values: 00, 06, 12, 18)"
        fi
        
        # 验证结束计数
        if ! [[ $end_count =~ ^[0-9]+$ ]]; then
            echo "Error: End count must be a number"
            exit 1
        fi
        
        echo "Starting download for single date..."
        download_gfs_files "$date_str" "$hour" "$end_count"
        local download_result=$?
        
        echo "Verifying download..."
        check_download_status "$date_str" "$hour" "$end_count"
        local check_result=$?
        
        if [ $check_result -eq 0 ]; then
            echo "Download completed successfully!"
            exit 0
        else
            echo "Download completed with $check_result missing/incomplete files"
            exit $check_result
        fi
    
    # 用法2: 下载日期范围数据
    # ./script.sh 20250701 20250930 12 168
    elif [[ $# -eq 4 ]]; then
        local start_date="$1"
        local end_date="$2"
        local hour="$3"
        local end_count="$4"
        
        echo "Starting batch download from $start_date to $end_date"
        generate_date_range "$start_date" "$end_date" "$hour" "$end_count"
    
    else
        echo "Usage:"
        echo "  Single date: $PROGNAME YYYYMMDD HH END_COUNT"
        echo "  Date range:  $PROGNAME START_DATE END_DATE HH END_COUNT"
        echo ""
        echo "Examples:"
        echo "  Download single date: $PROGNAME 20250701 12 168"
        echo "  Download date range:  $PROGNAME 20250701 20250930 12 168"
        exit 1
    fi
}

# 执行主函数
main "$@"
