#!/bin/bash
# ============================================================================
# MySQL 8.0.24 全量备份恢复脚本
# ============================================================================
# 作者: AI Assistant
# 版本: 1.0.0
# 说明: 从全量备份恢复 MySQL 数据库
# 
# 功能特点:
#   - 支持恢复全部数据库或指定数据库
#   - 自动解压压缩备份
#   - 恢复前可选择性备份当前数据
#   - 支持验证恢复结果
#
# 使用方法:
#   ./restore_full.sh [选项] <备份目录>
#
# 选项:
#   -d, --database <name>    仅恢复指定数据库
#   -y, --yes                跳过确认提示
#   --no-backup              恢复前不备份当前数据
#   -h, --help               显示帮助信息
#
# 警告:
#   恢复操作会覆盖现有数据，请确保已做好备份！
# ============================================================================

set -o pipefail

# ----------------------------------------------------------------------------
# 脚本路径和配置加载
# ----------------------------------------------------------------------------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 加载配置文件
if [[ ! -f "${SCRIPT_DIR}/config.sh" ]]; then
    echo "[ERROR] 配置文件不存在: ${SCRIPT_DIR}/config.sh"
    exit 1
fi
source "${SCRIPT_DIR}/config.sh"

# 加载公共函数库
if [[ ! -f "${SCRIPT_DIR}/lib/common.sh" ]]; then
    echo "[ERROR] 公共函数库不存在: ${SCRIPT_DIR}/lib/common.sh"
    exit 1
fi
source "${SCRIPT_DIR}/lib/common.sh"

# ----------------------------------------------------------------------------
# 全局变量
# ----------------------------------------------------------------------------
RESTORE_TIMESTAMP=$(date +"${TIMESTAMP_FORMAT}")
LOG_FILE="${LOG_DIR}/restore_full_${RESTORE_TIMESTAMP}.log"
LOCK_FILE="${BACKUP_ROOT_DIR}/.restore.lock"
BACKUP_PATH=""
SPECIFIED_DATABASE=""
SKIP_CONFIRM=false
NO_BACKUP_BEFORE_RESTORE=false

# ----------------------------------------------------------------------------
# 显示帮助信息
# ----------------------------------------------------------------------------
show_help() {
    cat << EOF
MySQL 8.0.24 全量备份恢复脚本

使用方法:
    $(basename "$0") [选项] <备份目录>

选项:
    -d, --database <name>    仅恢复指定数据库
    -y, --yes                跳过确认提示
    --no-backup              恢复前不备份当前数据
    -h, --help               显示此帮助信息

示例:
    $(basename "$0") /data/mysql_backup/full/mysql_backup_full_20231220_120000
    $(basename "$0") -d mydb /path/to/backup
    $(basename "$0") -y --no-backup /path/to/backup

警告:
    恢复操作会覆盖现有数据，请确保已做好备份！
EOF
}

# ----------------------------------------------------------------------------
# 解析命令行参数
# ----------------------------------------------------------------------------
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -d|--database)
                shift
                if [[ -z "$1" ]]; then
                    die "选项 -d/--database 需要一个参数"
                fi
                SPECIFIED_DATABASE="$1"
                shift
                ;;
            -y|--yes)
                SKIP_CONFIRM=true
                shift
                ;;
            --no-backup)
                NO_BACKUP_BEFORE_RESTORE=true
                shift
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            -*)
                die "未知选项: $1\n使用 --help 查看帮助"
                ;;
            *)
                if [[ -z "$BACKUP_PATH" ]]; then
                    BACKUP_PATH="$1"
                else
                    die "多余的参数: $1"
                fi
                shift
                ;;
        esac
    done
    
    if [[ -z "$BACKUP_PATH" ]]; then
        die "请指定备份目录\n使用 --help 查看帮助"
    fi
}

# ----------------------------------------------------------------------------
# 验证备份目录
# ----------------------------------------------------------------------------
validate_backup_dir() {
    log_info "验证备份目录: $BACKUP_PATH"
    
    if [[ ! -d "$BACKUP_PATH" ]]; then
        die "备份目录不存在: $BACKUP_PATH"
    fi
    
    # 检查是否有备份文件
    local sql_files=$(find "$BACKUP_PATH" -name "*.sql*" -type f 2>/dev/null)
    
    if [[ -z "$sql_files" ]]; then
        die "备份目录中没有找到 SQL 备份文件"
    fi
    
    # 检查元数据文件
    if [[ -f "${BACKUP_PATH}/metadata.txt" ]]; then
        log_info "备份元数据:"
        cat "${BACKUP_PATH}/metadata.txt" | grep -E "^(备份类型|备份时间|备份大小):" | while read line; do
            log_info "  $line"
        done
    fi
    
    log_info "备份目录验证通过"
}

# ----------------------------------------------------------------------------
# 显示备份内容并确认
# ----------------------------------------------------------------------------
confirm_restore() {
    echo ""
    echo "============================================================"
    echo "警告: 即将执行数据库恢复操作"
    echo "============================================================"
    echo ""
    echo "备份目录: $BACKUP_PATH"
    echo "目标服务器: ${MYSQL_HOST}:${MYSQL_PORT}"
    echo "MySQL 用户: $MYSQL_USER"
    
    if [[ -n "$SPECIFIED_DATABASE" ]]; then
        echo "恢复数据库: $SPECIFIED_DATABASE"
    else
        echo "恢复范围: 所有数据库"
    fi
    
    echo ""
    echo "备份文件列表:"
    find "$BACKUP_PATH" -name "*.sql*" -type f | while read f; do
        echo "  - $(basename "$f") ($(get_file_size "$f"))"
    done
    
    echo ""
    echo "警告: 此操作将覆盖目标数据库中的现有数据！"
    echo ""
    
    if [[ "$SKIP_CONFIRM" != "true" ]]; then
        read -p "确认要继续吗? (输入 'yes' 确认): " confirm
        if [[ "$confirm" != "yes" ]]; then
            log_info "操作已取消"
            exit 0
        fi
    else
        log_warn "跳过确认 (-y 参数)"
    fi
}

# ----------------------------------------------------------------------------
# 恢复前备份当前数据
# ----------------------------------------------------------------------------
backup_before_restore() {
    if [[ "$NO_BACKUP_BEFORE_RESTORE" == "true" ]]; then
        log_warn "跳过恢复前备份 (--no-backup 参数)"
        return
    fi
    
    log_info "恢复前备份当前数据..."
    
    local pre_restore_dir="${BACKUP_ROOT_DIR}/pre_restore_backup_${RESTORE_TIMESTAMP}"
    ensure_dir "$pre_restore_dir"
    
    local dump_args="-h${MYSQL_HOST} -P${MYSQL_PORT} -u${MYSQL_USER}"
    if [[ -n "$MYSQL_PASSWORD" ]]; then
        dump_args="$dump_args -p${MYSQL_PASSWORD}"
    fi
    
    if [[ -n "$SPECIFIED_DATABASE" ]]; then
        # 仅备份指定数据库
        log_info "备份数据库: $SPECIFIED_DATABASE"
        $MYSQLDUMP_PATH $dump_args "$SPECIFIED_DATABASE" > "${pre_restore_dir}/${SPECIFIED_DATABASE}.sql" 2>> "$LOG_FILE"
    else
        # 备份所有数据库
        log_info "备份所有数据库"
        $MYSQLDUMP_PATH $dump_args --all-databases > "${pre_restore_dir}/all_databases.sql" 2>> "$LOG_FILE"
    fi
    
    log_info "恢复前备份完成: $pre_restore_dir"
}

# ----------------------------------------------------------------------------
# 恢复单个数据库
# ----------------------------------------------------------------------------
restore_database() {
    local db_name="$1"
    local sql_file=""
    
    # 查找对应的备份文件
    sql_file=$(find "$BACKUP_PATH" -name "${db_name}.sql*" -type f 2>/dev/null | head -n1)
    
    if [[ -z "$sql_file" ]]; then
        die "找不到数据库 $db_name 的备份文件"
    fi
    
    log_info "恢复数据库: $db_name"
    log_info "备份文件: $sql_file"
    
    local mysql_args="-h${MYSQL_HOST} -P${MYSQL_PORT} -u${MYSQL_USER}"
    if [[ -n "$MYSQL_PASSWORD" ]]; then
        mysql_args="$mysql_args -p${MYSQL_PASSWORD}"
    fi
    
    # 判断是否需要解压
    if [[ "$sql_file" == *.gz ]]; then
        log_info "解压并恢复..."
        if ! gunzip -c "$sql_file" | $MYSQL_PATH $mysql_args "$db_name" 2>> "$LOG_FILE"; then
            die "数据库 $db_name 恢复失败"
        fi
    elif [[ "$sql_file" == *.lz4 ]]; then
        log_info "解压并恢复..."
        if ! lz4 -d -c "$sql_file" | $MYSQL_PATH $mysql_args "$db_name" 2>> "$LOG_FILE"; then
            die "数据库 $db_name 恢复失败"
        fi
    else
        log_info "直接恢复..."
        if ! $MYSQL_PATH $mysql_args "$db_name" < "$sql_file" 2>> "$LOG_FILE"; then
            die "数据库 $db_name 恢复失败"
        fi
    fi
    
    log_info "数据库 $db_name 恢复成功"
}

# ----------------------------------------------------------------------------
# 恢复所有数据库
# ----------------------------------------------------------------------------
restore_all_databases() {
    # 查找全库备份文件
    local all_db_file=$(find "$BACKUP_PATH" -name "all_databases.sql*" -type f 2>/dev/null | head -n1)
    
    if [[ -n "$all_db_file" ]]; then
        # 使用全库备份文件
        log_info "使用全库备份文件恢复"
        log_info "备份文件: $all_db_file"
        
        local mysql_args="-h${MYSQL_HOST} -P${MYSQL_PORT} -u${MYSQL_USER}"
        if [[ -n "$MYSQL_PASSWORD" ]]; then
            mysql_args="$mysql_args -p${MYSQL_PASSWORD}"
        fi
        
        if [[ "$all_db_file" == *.gz ]]; then
            log_info "解压并恢复..."
            if ! gunzip -c "$all_db_file" | $MYSQL_PATH $mysql_args 2>> "$LOG_FILE"; then
                die "全库恢复失败"
            fi
        elif [[ "$all_db_file" == *.lz4 ]]; then
            log_info "解压并恢复..."
            if ! lz4 -d -c "$all_db_file" | $MYSQL_PATH $mysql_args 2>> "$LOG_FILE"; then
                die "全库恢复失败"
            fi
        else
            log_info "直接恢复..."
            if ! $MYSQL_PATH $mysql_args < "$all_db_file" 2>> "$LOG_FILE"; then
                die "全库恢复失败"
            fi
        fi
        
        log_info "全库恢复成功"
    else
        # 逐个恢复数据库
        log_info "逐个恢复数据库"
        
        local sql_files=$(find "$BACKUP_PATH" -name "*.sql*" -type f | grep -v "all_databases")
        
        if [[ -z "$sql_files" ]]; then
            die "没有找到可恢复的数据库备份文件"
        fi
        
        while IFS= read -r sql_file; do
            local db_name=$(basename "$sql_file" | sed -E 's/\.sql(\.gz|\.lz4)?$//')
            
            # 先创建数据库 (如果不存在)
            log_info "确保数据库存在: $db_name"
            execute_mysql "CREATE DATABASE IF NOT EXISTS \`${db_name}\`;" 2>> "$LOG_FILE" || true
            
            restore_database "$db_name"
        done <<< "$sql_files"
    fi
}

# ----------------------------------------------------------------------------
# 验证恢复结果
# ----------------------------------------------------------------------------
verify_restore() {
    log_info "验证恢复结果..."
    
    if [[ -n "$SPECIFIED_DATABASE" ]]; then
        # 验证指定数据库
        local table_count=$(execute_mysql "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema='${SPECIFIED_DATABASE}';")
        log_info "数据库 $SPECIFIED_DATABASE 包含 $table_count 个表"
    else
        # 验证所有数据库
        local db_count=$(execute_mysql "SELECT COUNT(*) FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'performance_schema', 'sys', 'mysql');")
        log_info "共有 $db_count 个用户数据库"
    fi
    
    log_info "恢复验证完成"
}

# ----------------------------------------------------------------------------
# 初始化恢复环境
# ----------------------------------------------------------------------------
init_restore() {
    log_info "============================================================"
    log_info "MySQL 全量备份恢复开始"
    log_info "============================================================"
    log_info "恢复时间: $RESTORE_TIMESTAMP"
    log_info "备份目录: $BACKUP_PATH"
    log_info "日志文件: $LOG_FILE"
    log_info "============================================================"
    
    # 设置错误处理
    setup_error_trap
    
    # 检查必要命令
    check_commands "$MYSQL_PATH" "gunzip"
    
    # 检查 MySQL 连接
    check_mysql_connection
    
    # 创建必要目录
    ensure_dir "$LOG_DIR"
    ensure_dir "$RESTORE_TMP_DIR"
    
    # 获取锁
    acquire_lock "$LOCK_FILE" "恢复"
    
    # 设置清理 trap
    trap cleanup_restore EXIT
}

# ----------------------------------------------------------------------------
# 清理函数
# ----------------------------------------------------------------------------
cleanup_restore() {
    local exit_code=$?
    
    # 释放锁
    release_lock "$LOCK_FILE"
    
    if [[ $exit_code -ne 0 ]]; then
        log_error "恢复过程中发生错误，退出码: $exit_code"
        send_notification "[失败] MySQL 全量恢复" "恢复失败，请检查日志: $LOG_FILE"
    fi
}

# ----------------------------------------------------------------------------
# 主函数
# ----------------------------------------------------------------------------
main() {
    # 记录开始时间
    START_TIME=$(date +%s)
    
    # 解析参数
    parse_args "$@"
    
    # 初始化
    init_restore
    
    # 验证备份目录
    validate_backup_dir
    
    # 确认恢复操作
    confirm_restore
    
    # 恢复前备份
    backup_before_restore
    
    if [[ -n "$SPECIFIED_DATABASE" ]]; then
        # 恢复指定数据库
        # 先创建数据库 (如果不存在)
        log_info "确保数据库存在: $SPECIFIED_DATABASE"
        execute_mysql "CREATE DATABASE IF NOT EXISTS \`${SPECIFIED_DATABASE}\`;" 2>> "$LOG_FILE" || true
        
        restore_database "$SPECIFIED_DATABASE"
    else
        # 恢复所有数据库
        restore_all_databases
    fi
    
    # 验证恢复结果
    verify_restore
    
    # 计算总耗时
    local end_time=$(date +%s)
    local duration=$(calculate_duration "$START_TIME" "$end_time")
    
    log_info "============================================================"
    log_info "MySQL 全量备份恢复完成"
    log_info "============================================================"
    log_info "执行耗时: $duration"
    log_info "============================================================"
    
    # 发送成功通知
    send_notification "[成功] MySQL 全量恢复" "恢复完成\n耗时: $duration"
    
    return 0
}

# 执行主函数
main "$@"