#!/bin/bash
# ============================================================
# PubMed 基线下载+导入循环脚本
# 不断下载缺失/损坏文件 → 校验 → 导入，直到全部 138 个就绪
# 部署：scp 到服务器，setsid bash pubmed_download_loop.sh &
# ============================================================
set -e

DIR="${1:-/root/pubmed_baseline}"
COMPOSE_DIR="/root/scilit"
LOG_FILE="${DIR}/loop.log"
BASE_URL="https://ftp.ncbi.nlm.nih.gov/pubmed/baseline"

mkdir -p "$DIR"

echo "========================================================" | tee -a "$LOG_FILE"
echo "[$(date)] PubMed 基线下载循环启动" | tee -a "$LOG_FILE"
echo "目标目录: $DIR" | tee -a "$LOG_FILE"
echo "========================================================" | tee -a "$LOG_FILE"

cleanup() {
    echo "[$(date)] 收到退出信号，清理中..." | tee -a "$LOG_FILE"
    pkill -f "aria2c.*pubmed26n" 2>/dev/null || true
    exit 0
}
trap cleanup SIGTERM SIGINT

ROUND=0
while true; do
    ROUND=$((ROUND + 1))

    # ========== 检查哪些文件需要下载 ==========
    NEED_DOWNLOAD=()
    NEED_IMPORT=()

    for n in $(seq 1197 1334); do
        fname="pubmed26n$(printf '%04d' $n).xml.gz"
        fpath="$DIR/$fname"

        if [ -f "$fpath" ]; then
            # 有文件但可能是损坏的
            if gzip -t "$fpath" 2>/dev/null; then
                NEED_IMPORT+=("$fname")
            else
                echo "[$(date)]   [损坏] $fname，重新下载" | tee -a "$LOG_FILE"
                rm -f "$fpath"
                NEED_DOWNLOAD+=("$fname")
            fi
        else
            NEED_DOWNLOAD+=("$fname")
        fi
    done

    echo "[$(date)] ====== 第 $ROUND 轮 ======" | tee -a "$LOG_FILE"
    echo "[$(date)] 待下载: ${#NEED_DOWNLOAD[@]}  待导入: ${#NEED_IMPORT[@]}" | tee -a "$LOG_FILE"

    # ========== 导入已有文件 ==========
    if [ ${#NEED_IMPORT[@]} -gt 0 ]; then
        echo "[$(date)] --- 开始导入 ${#NEED_IMPORT[@]} 个文件 ---" | tee -a "$LOG_FILE"
        cd "$COMPOSE_DIR"
        docker compose -f docker-compose.prod.yml run --rm \
            -v "$DIR:/data/pubmed_baseline" \
            worker python /app/scripts/pubmed_baseline.py \
            --dir /data/pubmed_baseline 2>&1 | tee -a "$LOG_FILE"

        # 记录总量
        echo "[$(date)] 导入完成。DB 总量:" | tee -a "$LOG_FILE"
        docker exec scilit-postgres-1 psql -U scilit \
            -c "SELECT COUNT(*) FROM global_literature;" 2>/dev/null | tee -a "$LOG_FILE"

        # 导完的文件可以删掉（已入库）
        for fname in "${NEED_IMPORT[@]}"; do
            rm -f "$DIR/$fname"
        done
        echo "[$(date)] 已删除 ${#NEED_IMPORT[@]} 个已导入文件" | tee -a "$LOG_FILE"
    fi

    # ========== 检查是否全部完成 ==========
    if [ ${#NEED_DOWNLOAD[@]} -eq 0 ]; then
        echo "[$(date)] ====== 全部 138 个文件下载导入完成! ======" | tee -a "$LOG_FILE"
        # 汇总
        echo "[$(date)] DB 最终总量:" | tee -a "$LOG_FILE"
        docker exec scilit-postgres-1 psql -U scilit \
            -c "SELECT COUNT(*) FROM global_literature;" 2>/dev/null | tee -a "$LOG_FILE"
        exit 0
    fi

    # ========== 下载缺失文件 ==========
    echo "[$(date)] --- 开始下载 ${#NEED_DOWNLOAD[@]} 个文件 ---" | tee -a "$LOG_FILE"

    # 生成 URL 列表
    URL_FILE=$(mktemp)
    for fname in "${NEED_DOWNLOAD[@]}"; do
        echo "$BASE_URL/$fname" >> "$URL_FILE"
    done

    cd "$DIR"
    aria2c -x 5 -s 5 -k 1M --connect-timeout=30 --timeout=120 \
        --console-log-level=notice --summary-interval=30 \
        -i "$URL_FILE" 2>&1 | tee -a "$LOG_FILE"

    rm -f "$URL_FILE"

    echo "[$(date)] 下载完成。目录文件数: $(ls "$DIR"/*.xml.gz 2>/dev/null | wc -l)" | tee -a "$LOG_FILE"

    # 短暂停顿后进入下一轮
    sleep 5
done
