#!/bin/bash
# ============================================================
# 生产数据库终极升级脚本
# ============================================================
# 使用 session_replication_role 绕过外键约束
# 大幅提升升级速度和成功率
# ============================================================

set -e

# 配置
CONTAINER="cosmo_postgres"
DB_NAME="cosmo_db"
DB_USER="postgres"
BACKUP_FILE="backup_$(date +%Y%m%d_%H%M%S).sql"
SCRIPT_FILE="upgrade_production_final.sql"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 颜色
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

print_info() { echo -e "${BLUE}ℹ ${1}${NC}"; }
print_success() { echo -e "${GREEN}✅ ${1}${NC}"; }
print_warning() { echo -e "${YELLOW}⚠️  ${1}${NC}"; }
print_error() { echo -e "${RED}❌ ${1}${NC}"; }
print_step() { echo -e "${CYAN}▶ ${1}${NC}"; }

# 检查容器
check_container() {
    print_step "检查 Docker 容器状态..."
    if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER}$"; then
        print_error "容器 ${CONTAINER} 未运行！"
        docker ps --format "table {{.Names}}\t{{.Status}}"
        exit 1
    fi
    print_success "容器运行正常"
}

# 检查脚本
check_script() {
    print_step "检查升级脚本..."
    if [ ! -f "${SCRIPT_DIR}/${SCRIPT_FILE}" ]; then
        print_error "找不到 ${SCRIPT_FILE}"
        exit 1
    fi
    print_success "脚本就绪"
}

# 检查权限
check_permissions() {
    print_step "检查数据库权限..."

    SUPERUSER=$(docker exec ${CONTAINER} psql -U ${DB_USER} -d ${DB_NAME} -t -c \
        "SELECT usesuper FROM pg_user WHERE usename = current_user;" | tr -d ' ')

    if [ "$SUPERUSER" != "t" ]; then
        print_error "用户 ${DB_USER} 不是 superuser！"
        echo ""
        print_warning "session_replication_role 需要 superuser 权限"
        echo "解决方案："
        echo "  1. 使用 superuser 账号执行升级"
        echo "  2. 或临时授予权限: ALTER USER ${DB_USER} WITH SUPERUSER;"
        exit 1
    fi

    print_success "权限检查通过 (superuser)"
}

# 检查 display_name 字段
check_display_name() {
    print_step "检查 roles 表结构..."

    HAS_DISPLAY_NAME=$(docker exec ${CONTAINER} psql -U ${DB_USER} -d ${DB_NAME} -t -c \
        "SELECT COUNT(*) FROM information_schema.columns
         WHERE table_name = 'roles' AND column_name = 'display_name';" | tr -d ' ')

    if [ "$HAS_DISPLAY_NAME" = "1" ]; then
        print_info "检测到 display_name 字段（将使用对应版本）"
        echo ""
        print_warning "请确认 upgrade_production_final.sql 中："
        echo "  - 第 20-27 行（带 display_name）未注释"
        echo "  - 第 29-36 行（不带 display_name）已注释"
        echo ""
        read -p "是否确认脚本已正确配置? (yes/no): " confirm
        if [ "$confirm" != "yes" ]; then
            print_info "升级已取消，请检查脚本配置"
            exit 0
        fi
    else
        print_info "未检测到 display_name 字段"
        echo ""
        print_warning "请确认 upgrade_production_final.sql 中："
        echo "  - 第 20-27 行（带 display_name）已注释"
        echo "  - 第 29-36 行（不带 display_name）未注释"
        echo ""
        read -p "是否确认脚本已正确配置? (yes/no): " confirm
        if [ "$confirm" != "yes" ]; then
            print_info "升级已取消，请检查脚本配置"
            exit 0
        fi
    fi
}

# 备份数据库
backup_database() {
    print_step "备份数据库..."
    if docker exec ${CONTAINER} pg_dump -U ${DB_USER} -d ${DB_NAME} > "${SCRIPT_DIR}/${BACKUP_FILE}"; then
        SIZE=$(ls -lh "${SCRIPT_DIR}/${BACKUP_FILE}" | awk '{print $5}')
        print_success "备份完成: ${BACKUP_FILE} (${SIZE})"
    else
        print_error "备份失败！"
        exit 1
    fi
}

# 执行升级
execute_upgrade() {
    print_step "执行数据库升级..."
    echo "========================================================"

    if cat "${SCRIPT_DIR}/${SCRIPT_FILE}" | docker exec -i ${CONTAINER} psql -U ${DB_USER} -d ${DB_NAME}; then
        echo "========================================================"
        print_success "升级执行完成！"
        return 0
    else
        echo "========================================================"
        print_error "升级失败！"
        return 1
    fi
}

# 显示验证结果
show_verification() {
    print_step "数据验证..."
    echo ""

    docker exec ${CONTAINER} psql -U ${DB_USER} -d ${DB_NAME} -c "
    SELECT
        'celestial_bodies.short_name' as item,
        CASE WHEN EXISTS(
            SELECT 1 FROM information_schema.columns
            WHERE table_name='celestial_bodies' AND column_name='short_name'
        ) THEN '✓ 存在' ELSE '✗ 缺失' END as status
    UNION ALL
    SELECT
        'roles',
        COUNT(*)::text || ' 条记录'
    FROM roles
    UNION ALL
    SELECT
        'menus',
        COUNT(*)::text || ' 条记录'
    FROM menus
    UNION ALL
    SELECT
        'role_menus',
        COUNT(*)::text || ' 条记录'
    FROM role_menus
    UNION ALL
    SELECT
        'scheduled_jobs',
        COUNT(*)::text || ' 条记录'
    FROM scheduled_jobs
    UNION ALL
    SELECT
        'system_settings',
        COUNT(*)::text || ' 条记录'
    FROM system_settings;
    " -t

    echo ""
}

# 显示回滚信息
show_rollback_info() {
    echo ""
    print_warning "如需回滚，执行："
    echo "cat ${SCRIPT_DIR}/${BACKUP_FILE} | docker exec -i ${CONTAINER} psql -U ${DB_USER} -d ${DB_NAME}"
    echo ""
}

# 主函数
main() {
    echo "============================================================"
    echo "   生产数据库终极升级脚本"
    echo "   使用 session_replication_role 技术"
    echo "============================================================"
    echo ""

    # 检查
    check_container
    check_script
    check_permissions
    check_display_name

    # 确认
    echo ""
    print_warning "即将执行以下操作："
    echo "  1. 备份当前数据库"
    echo "  2. 使用 replica 模式绕过外键约束"
    echo "  3. 导入所有数据（无需关心顺序）"
    echo "  4. 恢复正常模式并验证数据完整性"
    echo ""
    echo "受影响的表："
    echo "  • celestial_bodies - 添加 short_name 字段"
    echo "  • roles - 创建/更新记录"
    echo "  • menus - 清空并重新导入 (14条)"
    echo "  • role_menus - 清空并重新导入 (16条)"
    echo "  • celestial_events - 清空"
    echo "  • scheduled_jobs - 清空并重新导入 (2条)"
    echo "  • system_settings - 导入/更新 (3条)"
    echo "  • user_roles - 为现有用户分配角色"
    echo ""

    read -p "是否继续? (yes/no): " confirm
    if [ "$confirm" != "yes" ]; then
        print_info "升级已取消"
        exit 0
    fi

    # 执行
    echo ""
    backup_database

    if execute_upgrade; then
        show_verification
        print_success "🎉 数据库升级成功！"
        show_rollback_info

        echo ""
        print_info "后续步骤："
        echo "  1. 重启后端服务: docker restart cosmo-backend"
        echo "  2. 登录系统验证菜单显示"
        echo "  3. 测试用户功能"
        echo ""
        exit 0
    else
        print_error "升级失败（已自动回滚）"
        show_rollback_info
        exit 1
    fi
}

main