#!/bin/bash
# ============================================================
# 数据库恢复脚本
# 
# 用法:
#   ./restore-db.sh --file /path/to/backup.sql.gz
#   ./restore-db.sh --latest                          # 恢复最近一次备份
#   ./restore-db.sh --date 2026-06-15                 # 恢复指定日期的备份
#   ./restore-db.sh --list                            # 列出所有可用备份
# ============================================================
set -euo pipefail

# ==================== 配置 ====================
DB_HOST="${DB_HOST:-localhost}"
DB_PORT="${DB_PORT:-5432}"
DB_NAME="${POSTGRES_DB:-water_management}"
DB_USER="${POSTGRES_USER:-water}"
DB_PASSWORD="${POSTGRES_PASSWORD:-}"
PGPASSWORD="${DB_PASSWORD}"
export PGPASSWORD

BACKUP_BASE="${BACKUP_DIR:-/opt/water-management/backups}"
BACKUP_DAILY="${BACKUP_BASE}/daily"
BACKUP_WEEKLY="${BACKUP_BASE}/weekly"
BACKUP_MONTHLY="${BACKUP_BASE}/monthly"

# ==================== 参数 ====================
BACKUP_FILE=""
MODE="file"  # file | latest | date | list
TARGET_DATE=""

while [[ $# -gt 0 ]]; do
    case $1 in
        --file)    BACKUP_FILE="$2"; MODE="file"; shift 2 ;;
        --latest)  MODE="latest"; shift ;;
        --date)    TARGET_DATE="$2"; MODE="date"; shift 2 ;;
        --list)    MODE="list"; shift ;;
        -h|--help)
            echo "用法: $0 [选项]"
            echo ""
            echo "选项:"
            echo "  --file PATH      指定备份文件路径"
            echo "  --latest         恢复最近一次备份"
            echo "  --date YYYY-MM-DD 恢复指定日期的备份"
            echo "  --list           列出所有可用备份"
            echo ""
            echo "示例:"
            echo "  $0 --file /opt/water-management/backups/daily/wm_water_management_20260615_020000.sql.gz"
            echo "  $0 --latest"
            echo "  $0 --date 2026-06-15"
            exit 0
            ;;
        *)
            echo "❌ 未知参数: $1"
            exit 1
            ;;
    esac
done

# ==================== 函数 ====================

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# 列出所有可用备份
list_backups() {
    echo "========================================="
    echo " 可用备份列表"
    echo "========================================="
    echo ""
    
    echo "📅 每日备份:"
    if [ -d "$BACKUP_DAILY" ]; then
        ls -lht "$BACKUP_DAILY"/*.sql.gz 2>/dev/null | head -10 || echo "  (无)"
    fi
    echo ""
    
    echo "📅 每周备份:"
    if [ -d "$BACKUP_WEEKLY" ]; then
        ls -lht "$BACKUP_WEEKLY"/*.sql.gz 2>/dev/null | head -10 || echo "  (无)"
    fi
    echo ""
    
    echo "📅 每月备份:"
    if [ -d "$BACKUP_MONTHLY" ]; then
        ls -lht "$BACKUP_MONTHLY"/*.sql.gz 2>/dev/null | head -10 || echo "  (无)"
    fi
    echo ""
}

# 查找最近的备份
find_latest_backup() {
    local latest=""
    latest=$(ls -1t "$BACKUP_DAILY"/*.sql.gz 2>/dev/null | head -1 || echo "")
    
    if [ -z "$latest" ]; then
        latest=$(ls -1t "$BACKUP_WEEKLY"/*.sql.gz 2>/dev/null | head -1 || echo "")
    fi
    
    if [ -z "$latest" ]; then
        latest=$(ls -1t "$BACKUP_MONTHLY"/*.sql.gz 2>/dev/null | head -1 || echo "")
    fi
    
    echo "$latest"
}

# 查找指定日期的备份
find_backup_by_date() {
    local date_str="$1"
    local date_compact
    date_compact=$(echo "$date_str" | tr -d '-')
    
    local found=""
    found=$(ls -1 "$BACKUP_DAILY"/*${date_compact}*.sql.gz 2>/dev/null | head -1 || echo "")
    
    echo "$found"
}

# 解密备份文件（如果需要）
decrypt_if_needed() {
    local file="$1"
    
    if [[ "$file" == *.enc ]]; then
        log "🔐 检测到加密文件，正在解密..."
        local decrypted="${file%.enc}"
        openssl enc -d -aes-256-cbc -pbkdf2 \
            -in "$file" \
            -out "$decrypted" \
            -pass "pass:${ENCRYPT_PASSPHRASE:-}"
        echo "$decrypted"
    else
        echo "$file"
    fi
}

# 恢复数据库
restore_database() {
    local backup_file="$1"
    
    if [ ! -f "$backup_file" ]; then
        log "❌ 备份文件不存在: $backup_file"
        exit 1
    fi
    
    local file_size
    file_size=$(du -sh "$backup_file" | cut -f1)
    
    log "========================================="
    log " 数据库恢复"
    log "========================================="
    log "备份文件: $backup_file"
    log "文件大小: $file_size"
    log "目标数据库: $DB_NAME"
    log "目标主机: $DB_HOST:$DB_PORT"
    log ""
    
    # 确认操作
    read -p "⚠️  此操作将覆盖当前数据库，是否继续？(yes/no): " confirm
    if [ "$confirm" != "yes" ]; then
        log "❌ 操作已取消"
        exit 0
    fi
    
    log "🔄 开始恢复..."
    
    # 解密（如果需要）
    local actual_file
    actual_file=$(decrypt_if_needed "$backup_file")
    
    # 恢复数据库
    if [[ "$actual_file" == *.gz ]]; then
        gunzip -c "$actual_file" | psql \
            -h "$DB_HOST" \
            -p "$DB_PORT" \
            -U "$DB_USER" \
            -d "$DB_NAME" \
            --single-transaction \
            -v ON_ERROR_STOP=1 \
            2>&1 | tail -20
    else
        psql \
            -h "$DB_HOST" \
            -p "$DB_PORT" \
            -U "$DB_USER" \
            -d "$DB_NAME" \
            --single-transaction \
            -v ON_ERROR_STOP=1 \
            -f "$actual_file" \
            2>&1 | tail -20
    fi
    
    if [ ${PIPESTATUS[0]} -eq 0 ]; then
        log "✅ 数据库恢复成功"
    else
        log "❌ 数据库恢复可能存在问题，请检查日志"
        exit 1
    fi
    
    # 清理解密的临时文件
    if [ "$actual_file" != "$backup_file" ] && [ -f "$actual_file" ]; then
        rm -f "$actual_file"
    fi
}

# ==================== 主逻辑 ====================

case "$MODE" in
    list)
        list_backups
        ;;
    latest)
        BACKUP_FILE=$(find_latest_backup)
        if [ -z "$BACKUP_FILE" ]; then
            log "❌ 未找到任何备份文件"
            exit 1
        fi
        restore_database "$BACKUP_FILE"
        ;;
    date)
        if [ -z "$TARGET_DATE" ]; then
            log "❌ 必须指定日期: --date YYYY-MM-DD"
            exit 1
        fi
        BACKUP_FILE=$(find_backup_by_date "$TARGET_DATE")
        if [ -z "$BACKUP_FILE" ]; then
            log "❌ 未找到 ${TARGET_DATE} 的备份文件"
            list_backups
            exit 1
        fi
        restore_database "$BACKUP_FILE"
        ;;
    file)
        if [ -z "$BACKUP_FILE" ]; then
            log "❌ 必须指定备份文件: --file /path/to/backup.sql.gz"
            exit 1
        fi
        restore_database "$BACKUP_FILE"
        ;;
    *)
        echo "❌ 未知模式: $MODE"
        exit 1
        ;;
esac