#!/bin/bash
# PubMed 基线 —— 单 aria2 实例 3 路并发 + 串行导入
# 用 aria2 的 -j 参数控制内部并发，比多进程更友好（共享连接池）
set -e

STATE_FILE="/tmp/pubmed_import_state.txt"
BATCH_DIR="/tmp/pubmed_batch"
SCRIPT="/app/scripts/pubmed_baseline.py"
BASE_URL="https://ftp.ncbi.nlm.nih.gov/pubmed/baseline"
TOTAL=1334
BATCH=3

mkdir -p "$BATCH_DIR"

CURRENT=$(cat "$STATE_FILE" 2>/dev/null || echo 0)
echo "[$(date)] Resuming from file $((CURRENT+1)), total $TOTAL"

for ((batch_start = CURRENT+1; batch_start <= TOTAL; batch_start += BATCH)); do
    batch_end=$((batch_start + BATCH - 1))
    if ((batch_end > TOTAL)); then batch_end=$TOTAL; fi

    echo "[$(date)] Batch $batch_start-$batch_end: downloading..."

    # 用单 aria2 实例下载本批，-j 限制并发数
    URLS=()
    for ((i = batch_start; i <= batch_end; i++)); do
        fnum=$(printf "%04d" $i)
        filename="pubmed26n${fnum}.xml.gz"
        filepath="$BATCH_DIR/$filename"

        if [ -f "$filepath" ] && [ -s "$filepath" ]; then
            echo "[$i] Already cached ($(stat -c%s "$filepath") bytes)"
            continue
        fi
        URLS+=("${BASE_URL}/pubmed26n${fnum}.xml.gz")
    done

    if [ ${#URLS[@]} -gt 0 ]; then
        aria2c -x 2 -s 2 -k 1M -j 3 --connect-timeout=30 --timeout=120 \
            --console-log-level=warn --summary-interval=0 \
            --retry-wait=5 --max-tries=3 \
            --dir="$BATCH_DIR" \
            "${URLS[@]}" 2>&1 | grep -E "^(Download Results| \[#|OK|ERR)"
    fi
    echo "[$(date)] Batch $batch_start-$batch_end: downloads complete"

    # 串行导入
    for ((i = batch_start; i <= batch_end; i++)); do
        fnum=$(printf "%04d" $i)
        filename="pubmed26n${fnum}.xml.gz"
        filepath="$BATCH_DIR/$filename"

        if [ ! -s "$filepath" ]; then
            echo "[$i] Download failed or empty, skipping"
            continue
        fi

        SIZE=$(stat -c%s "$filepath")
        if [ "$SIZE" -lt 1000000 ]; then
            echo "[$i] File too small ($SIZE bytes), skipping"
            rm -f "$filepath"
            continue
        fi

        echo "[$(date)] [$i] Importing $filename ($SIZE bytes)..."
        START=$(date +%s)

        docker exec scilit-backend-1 mkdir -p /tmp/pubmed_batch 2>/dev/null
        docker cp "$filepath" scilit-backend-1:/tmp/pubmed_batch/
        docker exec -w /app scilit-backend-1 python "$SCRIPT" --dir /tmp/pubmed_batch 2>&1 | tail -1

        END=$(date +%s)
        DURATION=$((END - START))

        echo "$i" > "$STATE_FILE"
        echo "[$i] ${DURATION}s" >> /tmp/pubmed_import.log
        echo "[$(date)] [$i] Done (${DURATION}s)"

        rm -f "$filepath"
        docker exec scilit-backend-1 rm -f "/tmp/pubmed_batch/$filename" 2>/dev/null
    done
done

echo "[$(date)] Done! All $TOTAL files processed."