#!/bin/bash

# 获取脚本所在目录的绝对路径
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)

# 切换工作目录到脚本所在目录
cd $SCRIPT_DIR

# 加载环境变量
load_env_variables() {
    if [ -f ".envrc" ]; then
        # 使用direnv加载环境变量
        if command -v direnv &> /dev/null; then
            eval "$(direnv export bash)"
            # 检查POSTGRES_PASSWORD是否已设置
            if [ -z "$POSTGRES_PASSWORD" ]; then
                echo "错误: 密码验证失败，无法继续操作"
                return 1
            fi
        else
            echo "错误: 未安装direnv，请先安装direnv"
            return 1
        fi
    else
        echo "错误: 找不到.envrc文件"
        return 1
    fi
    return 0
}

# 调用函数加载环境变量
if ! load_env_variables; then
    echo "无法加载环境变量，脚本退出"
    exit 1
fi

# 设置变量
container_name=${POSTGRES_CONTAINER_NAME:-postgres}
pg_user=${POSTGRES_USER:-postgres}
backup_dir="./data/backup/"

# 获取宿主机上的绝对备份路径
HOST_BACKUP_DIR="$(pwd)/${backup_dir}"

# 查找最新的备份文件
if [ -n "$1" ]; then
    # 如果指定了数据库名称，优先查找该数据库的最新备份
    target_db=$1
    echo "正在查找数据库 $target_db 的最新备份..."
    latest_backup=$(find "$backup_dir" -maxdepth 1 -type d -name "${target_db}_full_*" | sort -r | head -n 1)

    # 如果没有找到该数据库的备份，才考虑使用从备份文件名中提取的数据库名
    if [ -z "$latest_backup" ]; then
        echo "警告: 未找到数据库 $target_db 的备份，尝试使用从备份文件名中提取的数据库名查找"
        extracted_db_name=$(echo "$1" | sed -E 's/^([^_]+)_full_.+$|^([^_]+)$/\1\2/')
        latest_backup=$(find "$backup_dir" -maxdepth 1 -type d -name "${extracted_db_name}_full_*" | sort -r | head -n 1)
    fi
else
    # 如果未指定数据库名称，查找所有备份中的最新一个
    echo "未指定数据库名称，查找所有备份中的最新一个..."
    latest_backup=$(find "$backup_dir" -maxdepth 1 -type d -name "*_full_*" | sort -r | head -n 1)
fi

# 检查是否找到备份文件
if [ -z "$latest_backup" ]; then
    echo "错误: 在 $backup_dir 目录下未找到备份文件"
    exit 1
fi

# 提取备份文件名（不含路径）
backup_name=$(basename "$latest_backup")
HOST_BACKUP_PATH="${HOST_BACKUP_DIR}${backup_name}"

# 从备份文件名中提取数据库名（假设格式为 databaseName_full_timestamp）
extracted_db_name=$(echo "$backup_name" | sed -E 's/^([^_]+)_full_.+$/\1/')

# 目标数据库：优先使用命令行参数，否则使用从备份文件名中提取的数据库名
pg_database=${1:-$extracted_db_name}

# 验证找到的备份是否与目标数据库匹配
if [ -n "$1" ] && [ "$1" != "$extracted_db_name" ]; then
    echo "注意：找到的备份文件是 $backup_name，其中包含数据库 $extracted_db_name 的数据"
    echo "您指定的目标数据库是 $1，将把 $extracted_db_name 的数据恢复到 $1 数据库中"
    read -p "是否继续？(YES/no): " confirm_match
    if [ "$confirm_match" != "YES" ]; then
        echo "用户取消操作，脚本退出"
        exit 1
    fi
fi

# 显示找到的备份信息
echo "找到最新备份文件："
echo "备份目录名: $backup_name"
echo "宿主机上的备份路径: $HOST_BACKUP_PATH"

# 检查备份目录中是否包含toc.dat文件（确认是有效的PostgreSQL目录格式备份）
if [ ! -f "$latest_backup/toc.dat" ]; then
    echo "错误: 备份目录不包含toc.dat文件，可能不是有效的PostgreSQL目录格式备份"
    exit 1
fi

# 显示恢复目标信息
echo "\n恢复目标："
echo "目标数据库: $pg_database"
echo "目标容器: $container_name"

# 确认提示
echo "\n警告：此操作将恢复数据到数据库 $pg_database，可能会覆盖现有数据！"
read -p "是否继续执行恢复操作？(YES/no): " confirm
if [ "$confirm" != "YES" ]; then
    echo "用户取消恢复操作，脚本退出"
    exit 1
fi

# 检查数据库容器是否正在运行
if ! docker ps | grep -q "$container_name"; then
    echo "错误: 数据库容器 $container_name 未运行，请先启动服务"
    echo "您可以使用 ./service start 命令启动服务"
    exit 1
fi

# 检查目标数据库是否存在
if ! docker exec -i "$container_name" psql -U "$pg_user" -lqt | cut -d \| -f 1 | grep -qw "$pg_database"; then
    echo "警告: 目标数据库 $pg_database 不存在，将创建该数据库"
    if ! docker exec -i "$container_name" createdb -U "$pg_user" "$pg_database"; then
        echo "错误: 创建数据库 $pg_database 失败"
        exit 1
    fi
fi

# 执行恢复操作
echo "\n开始恢复数据库 $pg_database 从备份 $backup_name..."
echo "恢复过程可能需要一些时间，请耐心等待..."

# 使用pg_restore恢复数据库
if docker exec -i "$container_name" pg_restore \
    -U "$pg_user" \
    -d "$pg_database" \
    -Fd \
    -j 4 \
    "${backup_dir}${backup_name}"; then

    echo "\n数据库恢复成功!"
    echo "恢复详情："
    echo "- 备份源: $HOST_BACKUP_PATH"
    echo "- 目标数据库: $pg_database"
    echo "- 目标容器: $container_name"

    # 可选：显示数据库中的表数量，验证恢复结果
    table_count=$(docker exec -i "$container_name" psql -U "$pg_user" -d "$pg_database" -c "SELECT COUNT(*) FROM pg_tables WHERE schemaname NOT IN ('pg_catalog', 'information_schema');" -t -A)
    echo "- 恢复的表数量: $table_count"

    exit 0
else
    echo "\n数据库恢复失败!"
    echo "请检查错误信息并重试"
    exit 1
fi