#!/bin/bash
# 恢复基线导入：不断地重试失败的文件直到全部完成
set -euo pipefail

BATCH_DIR=/tmp/pubmed_batch
STATE_FILE=/tmp/pubmed_import_state.txt
BASE_URL=https://ftp.ncbi.nlm.nih.gov/pubmed/baseline
MIN_YEAR=2016
MAX_FILES=1334

mkdir -p "$BATCH_DIR"

echo "Starting baseline resume at $(date -u +'%Y-%m-%d %H:%M:%S UTC')"
echo "Target: files $(( $(cat "$STATE_FILE" 2>/dev/null || echo 0) + 1 )) - $MAX_FILES"

while true; do
    LAST_DONE=$(cat "$STATE_FILE" 2>/dev/null || echo 0)
    if [ "$LAST_DONE" -ge "$MAX_FILES" ]; then
        echo "All done! ($MAX_FILES files)"
        break
    fi

    i=$((LAST_DONE + 1))
    F=$(printf "pubmed26n%04d.xml.gz" $i)
    FILEPATH="$BATCH_DIR/$F"

    # 如果文件已存在且完整，跳过下载直接导入
    if [ -f "$FILEPATH" ] && gzip -t "$FILEPATH" 2>/dev/null; then
        echo "[$i] file exists, skip download"
    else
        rm -f "$FILEPATH" "${FILEPATH}.aria2"
        echo "[$i] downloading..."

        # aria2（多连接，快）
        OK=0
        for try in 1 2 3; do
            if aria2c -x 4 -s 4 --summary-interval=0 --console-log-level=error \
                --timeout=60 --connect-timeout=30 --retry-wait=10 \
                "$BASE_URL/$F" -d "$BATCH_DIR" -o "$F" 2>/dev/null; then
                OK=1; break
            fi
            sleep 15
        done

        # aria2 不行就 wget
        if [ "$OK" -ne 1 ]; then
            echo "[$i] aria2 failed, trying wget..."
            for try in 1 2 3; do
                rm -f "$FILEPATH"
                if wget -q --timeout=600 --tries=1 -O "$FILEPATH" "$BASE_URL/$F" 2>/dev/null; then
                    OK=1; break
                fi
                sleep 30
            done
        fi

        if [ "$OK" -ne 1 ]; then
            echo "[$i] all downloads failed, waiting 5 min then retry"
            sleep 300
            continue
        fi

        # 校验
        if ! gzip -t "$FILEPATH" 2>/dev/null; then
            echo "[$i] CRC fail, redownload"
            rm -f "$FILEPATH"
            continue
        fi
    fi

    START=$(date +%s)

    # 导入
    docker compose -f /root/scilit/docker-compose.prod.yml run --rm --no-deps \
        -v /root/scilit/backend/scripts/pubmed_baseline.py:/app/scripts/pubmed_baseline.py \
        -v "$BATCH_DIR:/tmp/pubmed_batch" \
        backend python /app/scripts/pubmed_baseline.py \
        --dir /tmp/pubmed_batch --min-year "$MIN_YEAR" --max-files 1 > /dev/null 2>&1
    EC=$?
    ELAPSED=$(( $(date +%s) - START ))

    rm -f "$FILEPATH" "${FILEPATH}.aria2"

    if [ $EC -ne 0 ]; then
        echo "[$i] import ERROR (${ELAPSED}s), will retry"
        sleep 10
        continue  # not mark as done
    fi

    echo "$i" > "$STATE_FILE"
    echo "[$i] OK (${ELAPSED}s)"
    sleep 3
done