#!/bin/bash
# PubMed 基线 —— aria2 下载 + 流水线（下载/导入并行）
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

mkdir -p "$BATCH_DIR"

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

aria2c --version | head -1

for ((i = CURRENT+1; i <= TOTAL; i++)); do
    fnum=$(printf "%04d" $i)
    filename="pubmed26n${fnum}.xml.gz"
    filepath="$BATCH_DIR/$filename"

    # Download if not cached
    if [ ! -f "$filepath" ]; then
        echo "[$(date)] [$i] Downloading $filename ..."
        aria2c -x 5 -s 5 -k 1M --connect-timeout=30 --timeout=60 \
            --console-log-level=warn --summary-interval=0 \
            -d "$BATCH_DIR" -o "$filename" \
            "${BASE_URL}/${filename}" 2>&1 | tail -3

        if [ ! -s "$filepath" ]; then
            echo "[$(date)] [$i] Retrying..."
            sleep 3
            aria2c -x 5 -s 5 -k 1M --connect-timeout=30 --timeout=60 \
                --console-log-level=warn --summary-interval=0 \
                -d "$BATCH_DIR" -o "$filename" \
                "${BASE_URL}/${filename}" 2>&1 | tail -3
        fi

        if [ ! -s "$filepath" ]; then
            echo "[$(date)] [$i] Download failed, skipping"
            rm -f "$filepath"
            continue
        fi
    fi

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

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

    # Copy to container and import
    docker cp "$filepath" scilit-backend-1:/tmp/pubmed_batch/ 2>/dev/null
    docker exec -w /app scilit-backend-1 python "$SCRIPT" --dir /tmp/pubmed_batch 2>&1 | tail -1

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

    # Update state and log
    echo "$i" > "$STATE_FILE"
    echo "[$i] ${DURATION}s" >> /tmp/pubmed_import.log

    echo "[$(date)] [$i] Done (${DURATION}s)"

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

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