# ================================================
# AI临床研究平台 - 数据库迁移脚本
# 日期: 2026-01-26
# 功能: 
#   1. RDS数据库重命名 (ai_clinical_research -> ai_clinical_research_test)
#   2. 安装pg_bigm和pgvector插件
#   3. 本地数据库导出
#   4. 导入到RDS测试数据库
# ================================================

Write-Host "========================================" -ForegroundColor Green
Write-Host " AI临床研究平台 - 数据库迁移脚本" -ForegroundColor Green
Write-Host " 执行时间: $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')" -ForegroundColor Green
Write-Host "========================================" -ForegroundColor Green

# ============ 配置区域 ============
$RDS_HOST_EXTERNAL = "pgm-2zex1m2y3r23hdn5ko.pg.rds.aliyuncs.com"  # RDS外网地址（需要先开启）
$RDS_HOST_INTERNAL = "pgm-2zex1m2y3r23hdn5.pg.rds.aliyuncs.com"    # RDS内网地址
$RDS_PORT = "5432"
$RDS_USER = "airesearch"
$RDS_PASSWORD = "Xibahe@fengzhibo117"
$OLD_DB_NAME = "ai_clinical_research"
$NEW_DB_NAME = "ai_clinical_research_test"

$LOCAL_CONTAINER = "ai-clinical-postgres"
$LOCAL_USER = "postgres"
$LOCAL_DB = "ai_clinical_research"

$BACKUP_FILE = "local_db_export_$(Get-Date -Format 'yyyyMMdd_HHmmss').sql"

# ============ Step 0: 前置检查 ============
Write-Host "`n[Step 0] 前置检查..." -ForegroundColor Cyan

# 检查Docker容器
$containerStatus = docker ps --filter "name=$LOCAL_CONTAINER" --format "{{.Status}}"
if (-not $containerStatus) {
    Write-Host "❌ 本地Docker容器 $LOCAL_CONTAINER 未运行！" -ForegroundColor Red
    exit 1
}
Write-Host "✅ 本地Docker容器运行中: $containerStatus" -ForegroundColor Green

# ============ Step 1: 导出本地数据库 ============
Write-Host "`n[Step 1] 导出本地数据库..." -ForegroundColor Cyan
Write-Host "   文件: $BACKUP_FILE" -ForegroundColor Yellow

docker exec $LOCAL_CONTAINER pg_dump `
    -U $LOCAL_USER `
    -d $LOCAL_DB `
    --format=plain `
    --no-owner `
    --no-acl `
    --encoding=UTF8 `
    > $BACKUP_FILE

if ($LASTEXITCODE -ne 0) {
    Write-Host "❌ 数据库导出失败！" -ForegroundColor Red
    exit 1
}

$fileSize = (Get-Item $BACKUP_FILE).Length / 1MB
Write-Host "✅ 数据库导出成功！文件大小: $([math]::Round($fileSize, 2)) MB" -ForegroundColor Green

# ============ Step 2: 连接RDS并重命名数据库 ============
Write-Host "`n[Step 2] 连接RDS并重命名数据库..." -ForegroundColor Cyan
Write-Host "   注意: 需要先在阿里云控制台开启RDS外网访问！" -ForegroundColor Yellow
Write-Host "   外网地址: $RDS_HOST_EXTERNAL" -ForegroundColor Yellow

$env:PGPASSWORD = $RDS_PASSWORD

# 2.1 检查RDS连接
Write-Host "`n   2.1 测试RDS连接..." -ForegroundColor Yellow
$testResult = docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL `
    -p $RDS_PORT `
    -U $RDS_USER `
    -d postgres `
    -c "SELECT version();" 2>&1

if ($LASTEXITCODE -ne 0) {
    Write-Host "❌ 无法连接RDS！请确保：" -ForegroundColor Red
    Write-Host "   1. RDS外网地址已开启" -ForegroundColor Red
    Write-Host "   2. 白名单已配置 0.0.0.0/0（临时）" -ForegroundColor Red
    Write-Host "   错误信息: $testResult" -ForegroundColor Red
    exit 1
}
Write-Host "✅ RDS连接成功！" -ForegroundColor Green

# 2.2 检查旧数据库是否存在
Write-Host "`n   2.2 检查数据库 $OLD_DB_NAME 是否存在..." -ForegroundColor Yellow
$dbExists = docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL `
    -p $RDS_PORT `
    -U $RDS_USER `
    -d postgres `
    -t -c "SELECT 1 FROM pg_database WHERE datname = '$OLD_DB_NAME';"

if ($dbExists.Trim() -eq "1") {
    Write-Host "✅ 数据库 $OLD_DB_NAME 存在，准备重命名..." -ForegroundColor Green
    
    # 2.3 断开所有连接
    Write-Host "`n   2.3 断开所有活动连接..." -ForegroundColor Yellow
    docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
        -h $RDS_HOST_EXTERNAL `
        -p $RDS_PORT `
        -U $RDS_USER `
        -d postgres `
        -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$OLD_DB_NAME' AND pid <> pg_backend_pid();"
    
    # 2.4 重命名数据库
    Write-Host "`n   2.4 重命名数据库 $OLD_DB_NAME -> $NEW_DB_NAME ..." -ForegroundColor Yellow
    docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
        -h $RDS_HOST_EXTERNAL `
        -p $RDS_PORT `
        -U $RDS_USER `
        -d postgres `
        -c "ALTER DATABASE $OLD_DB_NAME RENAME TO $NEW_DB_NAME;"
    
    if ($LASTEXITCODE -eq 0) {
        Write-Host "✅ 数据库重命名成功！" -ForegroundColor Green
    } else {
        Write-Host "❌ 数据库重命名失败！" -ForegroundColor Red
        exit 1
    }
} else {
    Write-Host "⚠️ 数据库 $OLD_DB_NAME 不存在，将直接创建 $NEW_DB_NAME" -ForegroundColor Yellow
}

# ============ Step 3: 创建/重建测试数据库 ============
Write-Host "`n[Step 3] 准备测试数据库 $NEW_DB_NAME ..." -ForegroundColor Cyan

# 3.1 检查测试数据库是否存在
$testDbExists = docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL `
    -p $RDS_PORT `
    -U $RDS_USER `
    -d postgres `
    -t -c "SELECT 1 FROM pg_database WHERE datname = '$NEW_DB_NAME';"

if ($testDbExists.Trim() -eq "1") {
    Write-Host "   数据库 $NEW_DB_NAME 已存在" -ForegroundColor Yellow
    
    # 询问是否清空重建
    $confirm = Read-Host "   是否删除现有数据并重新导入? (y/n)"
    if ($confirm -eq "y") {
        Write-Host "   断开连接并删除数据库..." -ForegroundColor Yellow
        docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
            -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d postgres `
            -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$NEW_DB_NAME';"
        docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
            -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d postgres `
            -c "DROP DATABASE $NEW_DB_NAME;"
        
        # 重新创建
        Write-Host "   创建数据库 $NEW_DB_NAME ..." -ForegroundColor Yellow
        docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
            -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d postgres `
            -c "CREATE DATABASE $NEW_DB_NAME WITH ENCODING='UTF8' LC_COLLATE='en_US.utf8' LC_CTYPE='en_US.utf8' TEMPLATE=template0;"
    }
} else {
    Write-Host "   创建数据库 $NEW_DB_NAME ..." -ForegroundColor Yellow
    docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
        -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d postgres `
        -c "CREATE DATABASE $NEW_DB_NAME WITH ENCODING='UTF8' LC_COLLATE='en_US.utf8' LC_CTYPE='en_US.utf8' TEMPLATE=template0;"
}

Write-Host "✅ 测试数据库准备完成！" -ForegroundColor Green

# ============ Step 4: 安装PostgreSQL插件 ============
Write-Host "`n[Step 4] 安装PostgreSQL插件..." -ForegroundColor Cyan

# 4.1 安装pg_bigm
Write-Host "   4.1 安装 pg_bigm 插件..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "CREATE EXTENSION IF NOT EXISTS pg_bigm;"

# 4.2 安装pgvector
Write-Host "   4.2 安装 pgvector 插件..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "CREATE EXTENSION IF NOT EXISTS vector;"

# 4.3 验证插件安装
Write-Host "`n   4.3 验证插件安装..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "SELECT extname, extversion FROM pg_extension WHERE extname IN ('pg_bigm', 'vector');"

Write-Host "✅ 插件安装完成！" -ForegroundColor Green

# ============ Step 5: 导入数据到RDS ============
Write-Host "`n[Step 5] 导入数据到RDS测试数据库..." -ForegroundColor Cyan
Write-Host "   这可能需要几分钟，请耐心等待..." -ForegroundColor Yellow

Get-Content $BACKUP_FILE | docker exec -i -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL `
    -p $RDS_PORT `
    -U $RDS_USER `
    -d $NEW_DB_NAME

if ($LASTEXITCODE -eq 0) {
    Write-Host "✅ 数据导入成功！" -ForegroundColor Green
} else {
    Write-Host "⚠️ 数据导入过程中可能有警告，请检查日志" -ForegroundColor Yellow
}

# ============ Step 6: 验证导入结果 ============
Write-Host "`n[Step 6] 验证导入结果..." -ForegroundColor Cyan

# 6.1 检查Schema数量
Write-Host "   6.1 检查Schema列表..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "SELECT nspname FROM pg_namespace WHERE nspname LIKE '%_schema' ORDER BY nspname;"

# 6.2 检查表数量
Write-Host "`n   6.2 检查各Schema表数量..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "SELECT schemaname, COUNT(*) as table_count FROM pg_tables WHERE schemaname NOT IN ('pg_catalog', 'information_schema') GROUP BY schemaname ORDER BY schemaname;"

# 6.3 检查iit_schema的表
Write-Host "`n   6.3 检查 iit_schema 的表..." -ForegroundColor Yellow
docker exec -e PGPASSWORD=$RDS_PASSWORD $LOCAL_CONTAINER psql `
    -h $RDS_HOST_EXTERNAL -p $RDS_PORT -U $RDS_USER -d $NEW_DB_NAME `
    -c "SELECT tablename FROM pg_tables WHERE schemaname = 'iit_schema';"

# ============ Step 7: 提醒关闭外网访问 ============
Write-Host "`n========================================" -ForegroundColor Red
Write-Host " ⚠️  重要安全提醒  ⚠️" -ForegroundColor Red
Write-Host "========================================" -ForegroundColor Red
Write-Host "数据库迁移完成！" -ForegroundColor Green
Write-Host ""
Write-Host "请立即执行以下安全操作：" -ForegroundColor Yellow
Write-Host "1. 登录阿里云RDS控制台" -ForegroundColor Yellow
Write-Host "2. 释放RDS外网地址" -ForegroundColor Yellow
Write-Host "3. 恢复白名单为 172.17.0.0/16（移除 0.0.0.0/0）" -ForegroundColor Yellow
Write-Host ""
Write-Host "新的测试数据库连接字符串（内网）：" -ForegroundColor Cyan
Write-Host "DATABASE_URL=postgresql://${RDS_USER}:Xibahe%40fengzhibo117@${RDS_HOST_INTERNAL}:5432/${NEW_DB_NAME}?connection_limit=18&pool_timeout=10" -ForegroundColor Cyan
Write-Host ""
Write-Host "备份文件: $BACKUP_FILE" -ForegroundColor Yellow
Write-Host "========================================" -ForegroundColor Green