#!/bin/bash
# ============================================================
# PubMed 基线下载后自动导入监控脚本
# 检测 aria2 下载完成后自动执行 docker 导入
# 部署：scp 到生产服务器后 nohup bash auto_import.sh &
# ============================================================
set -e

DATA_DIR="${1:-/root/pubmed_baseline}"
COMPOSE_DIR="/root/scilit"
LOG_FILE="/root/auto_import.log"
PID_FILE="/tmp/auto_import.pid"

echo "[$(date)] Auto-import monitor started, watching: $DATA_DIR" | tee -a "$LOG_FILE"
echo $$ > "$PID_FILE"

cleanup() {
    echo "[$(date)] Monitor stopping..." | tee -a "$LOG_FILE"
    rm -f "$PID_FILE"
    exit 0
}
trap cleanup SIGTERM SIGINT

while true; do
    # 检查 aria2 进程是否还在运行
    aria2_count=$(ps aux | grep a[r]ia2c | wc -l)

    if [ "$aria2_count" -eq 0 ]; then
        # aria2 已退出，检查是否有文件待导入
        file_count=$(ls "$DATA_DIR"/*.xml.gz 2>/dev/null | wc -l)
        if [ "$file_count" -gt 0 ]; then
            echo "[$(date)] aria2 done, $file_count files to import. Starting import..." | tee -a "$LOG_FILE"

            # 先校验文件完整性
            bad_files=0
            for f in "$DATA_DIR"/*.xml.gz; do
                if ! gzip -t "$f" 2>/dev/null; then
                    echo "[$(date)] CORRUPT: $f, removing and will re-download" | tee -a "$LOG_FILE"
                    rm -f "$f"
                    bad_files=$((bad_files + 1))
                fi
            done
            echo "[$(date)] Integrity check done: $bad_files corrupt files removed" | tee -a "$LOG_FILE"

            # 检查是否仍有文件需要导入
            remaining=$(ls "$DATA_DIR"/*.xml.gz 2>/dev/null | wc -l)
            if [ "$remaining" -eq 0 ]; then
                echo "[$(date)] No files to import after integrity check" | tee -a "$LOG_FILE"
                exit 1
            fi

            # 执行导入
            echo "[$(date)] Importing $remaining files..." | tee -a "$LOG_FILE"
            cd "$COMPOSE_DIR"
            docker compose -f docker-compose.prod.yml run --rm \
                -v "$DATA_DIR:/data/pubmed_baseline" \
                worker python /app/scripts/pubmed_baseline.py \
                --dir /data/pubmed_baseline 2>&1 | tee -a "$LOG_FILE"

            EXIT_CODE=${PIPESTATUS[0]}
            echo "[$(date)] Import exit code: $EXIT_CODE" | tee -a "$LOG_FILE"

            # 记录单个文件导入结果
            echo "[$(date)] Import completed. DB count:" | 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"

            break
        fi
    fi

    # aria2 仍在运行，等待
    sleep 30
done

echo "[$(date)] All done." | tee -a "$LOG_FILE"
