#!/bin/bash

# QAUP 数据库管理脚本
# 统一管理数据库相关的所有操作：导出、导入、备份、恢复、健康检查等

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# 默认配置
DB_HOST=${DB_HOST:-localhost}
DB_PORT=${DB_PORT:-5432}
DB_NAME=${DB_NAME:-qaup}
DB_USER=${DB_USER:-postgres}
EXPORT_DIR=${EXPORT_DIR:-$SCRIPT_DIR/export}

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

print_message() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
}

# 显示帮助信息
show_help() {
    echo "QAUP 数据库管理脚本"
    echo ""
    echo "用法: $0 <命令> [选项]"
    echo ""
    echo "数据管理命令:"
    echo "  prepare-default     准备默认初始数据"
    echo "  prepare-export      从现有数据库导出并准备数据"
    echo "  prepare-custom      自定义数据组合"
    echo "  prepare-status      显示数据文件状态"
    echo "  prepare-clean       清理数据文件"
    echo ""
    echo "导出命令:"
    echo "  export-all          导出完整数据库"
    echo "  export-schema       仅导出表结构"
    echo "  export-data         仅导出数据"
    echo "  export-initial      导出分层初始数据（推荐）"
    echo ""
    echo "备份恢复命令:"
    echo "  backup              备份数据库"
    echo "  restore <file>      恢复数据库"
    echo ""
    echo "维护命令:"
    echo "  health              健康检查"
    echo "  test-connection     测试数据库连接"
    echo ""
    echo "环境变量:"
    echo "  DB_HOST            数据库主机（默认：localhost）"
    echo "  DB_PORT            数据库端口（默认：5432）"
    echo "  DB_NAME            数据库名称（默认：qaup）"
    echo "  DB_USER            数据库用户（默认：postgres）"
    echo "  EXPORT_DIR         导出目录（默认：./export）"
    echo ""
    echo "示例:"
    echo "  $0 prepare-default              # 准备默认数据"
    echo "  $0 prepare-export               # 从现有系统导出数据"
    echo "  $0 export-initial               # 导出分层初始数据"
    echo "  $0 backup                       # 备份数据库"
    echo "  $0 health                       # 健康检查"
}

# ==========================================
# 数据准备功能
# ==========================================

# 显示数据文件状态
show_data_status() {
    print_message $BLUE "当前数据文件状态:"
    echo ""
    
    local files=(
        "initial_data.sql:实际使用的初始数据"
        "initial_data.template.sql:默认数据模板"
        "01_system_data.sql:系统配置数据"
        "02_business_data.sql:业务基础数据"
        "03_sample_data.sql:示例数据"
    )
    
    for file_info in "${files[@]}"; do
        local file="${file_info%%:*}"
        local desc="${file_info##*:}"
        local path="$SCRIPT_DIR/$file"
        
        if [ -f "$path" ]; then
            local size=$(du -h "$path" | cut -f1)
            local lines=$(wc -l < "$path")
            printf "  %-25s %s %8s %6s 行 - %s\n" "$file" "✓" "$size" "$lines" "$desc"
        else
            printf "  %-25s %s %8s %6s    - %s\n" "$file" "✗" "-" "-" "$desc"
        fi
    done
    
    echo ""
    
    if [ -f "$SCRIPT_DIR/initial_data.sql" ]; then
        print_message $GREEN "部署时将使用: initial_data.sql"
    else
        print_message $YELLOW "部署时将按顺序使用分层数据文件"
    fi
}

# 准备默认数据
prepare_default_data() {
    print_message $BLUE "准备默认模板数据..."
    
    if [ ! -f "$SCRIPT_DIR/initial_data.template.sql" ]; then
        print_message $RED "错误: 默认模板文件不存在"
        exit 1
    fi
    
    cp "$SCRIPT_DIR/initial_data.template.sql" "$SCRIPT_DIR/initial_data.sql"
    
    print_message $GREEN "✓ 已准备默认数据"
    print_message $YELLOW "默认管理员账号: admin"
    print_message $YELLOW "默认管理员密码: admin123"
}

# 从现有数据库导出并准备数据
prepare_export_data() {
    print_message $BLUE "从现有数据库导出数据..."
    
    # 检查连接参数
    if [ -z "$DB_HOST" ] || [ -z "$DB_NAME" ] || [ -z "$DB_USER" ]; then
        print_message $YELLOW "请设置数据库连接环境变量或使用默认值"
        echo "当前设置:"
        echo "  DB_HOST=$DB_HOST"
        echo "  DB_PORT=$DB_PORT"
        echo "  DB_NAME=$DB_NAME"
        echo "  DB_USER=$DB_USER"
        echo ""
        read -p "是否继续? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi
    
    # 测试连接
    test_database_connection
    
    # 导出分层数据
    export_initial_data
    
    # 使用导出的完整数据
    if [ -f "$EXPORT_DIR/initial_data_complete.sql" ]; then
        cp "$EXPORT_DIR/initial_data_complete.sql" "$SCRIPT_DIR/initial_data.sql"
        print_message $GREEN "✓ 已导出并准备完整初始数据"
    else
        print_message $RED "错误: 导出失败"
        exit 1
    fi
}

# 自定义数据组合
prepare_custom_data() {
    print_message $BLUE "自定义数据组合..."
    
    local system_data="$SCRIPT_DIR/01_system_data.sql"
    local business_data="$SCRIPT_DIR/02_business_data.sql"
    local sample_data="$SCRIPT_DIR/03_sample_data.sql"
    
    if [ ! -f "$system_data" ]; then
        print_message $RED "错误: 未找到系统数据文件，请先导出数据"
        print_message $YELLOW "提示: $0 export-initial"
        exit 1
    fi
    
    echo "请选择要包含的数据类型:"
    echo "1. 系统配置数据 (必需) ✓"
    echo "2. 业务基础数据"
    echo "3. 示例数据"
    echo ""
    
    local include_business=false
    local include_sample=false
    
    if [ -f "$business_data" ]; then
        read -p "是否包含业务基础数据? (Y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            include_business=true
        fi
    fi
    
    if [ -f "$sample_data" ]; then
        read -p "是否包含示例数据? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            include_sample=true
        fi
    fi
    
    # 生成自定义数据文件
    local output_file="$SCRIPT_DIR/initial_data.sql"
    
    cat > "$output_file" << 'EOF'
-- QAUP 自定义初始数据组合
-- 此文件由 db-manager.sh 自动生成

EOF
    
    # 添加系统数据
    echo "-- 系统基础配置数据" >> "$output_file"
    cat "$system_data" >> "$output_file"
    print_message $GREEN "✓ 已包含系统配置数据"
    
    # 添加业务数据
    if [ "$include_business" = true ] && [ -f "$business_data" ]; then
        echo "-- 业务基础数据" >> "$output_file"
        cat "$business_data" >> "$output_file"
        print_message $GREEN "✓ 已包含业务基础数据"
    fi
    
    # 添加示例数据
    if [ "$include_sample" = true ] && [ -f "$sample_data" ]; then
        echo "-- 示例数据" >> "$output_file"
        cat "$sample_data" >> "$output_file"
        print_message $GREEN "✓ 已包含示例数据"
    fi
    
    print_message $GREEN "✓ 自定义数据组合完成"
}

# 清理数据文件
clean_data_files() {
    print_message $YELLOW "清理数据文件..."
    
    read -p "确定要清理所有数据文件吗? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        return
    fi
    
    local files_to_clean=(
        "initial_data.sql"
        "01_system_data.sql"
        "02_business_data.sql"
        "03_sample_data.sql"
    )
    
    for file in "${files_to_clean[@]}"; do
        local path="$SCRIPT_DIR/$file"
        if [ -f "$path" ]; then
            rm "$path"
            print_message $GREEN "✓ 已删除 $file"
        fi
    done
    
    if [ -d "$EXPORT_DIR" ]; then
        rm -rf "$EXPORT_DIR"
        print_message $GREEN "✓ 已清理导出目录"
    fi
}

# ==========================================
# 数据导出功能
# ==========================================

# 测试数据库连接
test_database_connection() {
    print_message $BLUE "测试数据库连接..."
    
    if command -v pg_isready &> /dev/null; then
        if pg_isready -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" &>/dev/null; then
            print_message $GREEN "✓ 数据库连接成功"
        else
            print_message $RED "✗ 数据库连接失败"
            exit 1
        fi
    else
        print_message $YELLOW "警告: pg_isready 不可用，跳过连接测试"
    fi
}

# 导出完整数据库
export_all_data() {
    print_message $BLUE "导出完整数据库..."
    
    test_database_connection
    mkdir -p "$EXPORT_DIR"
    
    local output_file="$EXPORT_DIR/qaup_complete_backup.sql"
    
    pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        --no-owner --no-privileges --clean --if-exists \
        > "$output_file"
    
    print_message $GREEN "✓ 完整数据库导出完成: $output_file"
}

# 导出表结构
export_schema_only() {
    print_message $BLUE "导出数据库表结构..."
    
    test_database_connection
    mkdir -p "$EXPORT_DIR"
    
    local output_file="$EXPORT_DIR/qaup_database_schema.sql"
    
    pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        --schema-only --no-owner --no-privileges --clean --if-exists \
        > "$output_file"
    
    print_message $GREEN "✓ 表结构导出完成: $output_file"
}

# 导出分层初始数据
export_initial_data() {
    print_message $BLUE "导出分层初始数据..."
    
    test_database_connection
    mkdir -p "$EXPORT_DIR"
    
    # 导出系统基础数据
    export_system_data
    
    # 导出业务基础数据
    export_business_data
    
    # 导出示例数据
    export_sample_data
    
    # 生成合并文件
    generate_combined_data
}

# 导出系统基础数据
export_system_data() {
    local output_file="$EXPORT_DIR/01_system_data.sql"
    local system_tables=(
        "sys_config" "sys_dict_type" "sys_dict_data" "sys_dept" "sys_post"
        "sys_role" "sys_menu" "sys_user" "sys_user_role" "sys_user_post"
        "sys_role_menu" "sys_role_dept" "sys_notice"
    )
    
    local table_args=""
    for table in "${system_tables[@]}"; do
        table_args="$table_args --table=$table"
    done
    
    cat > "$output_file" << 'EOF'
-- QAUP 系统基础配置数据
SET session_replication_role = replica;
EOF
    
    pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        --data-only --no-owner --no-privileges --disable-triggers --inserts \
        $table_args >> "$output_file"
    
    cat >> "$output_file" << 'EOF'
SET session_replication_role = DEFAULT;
EOF
    
    print_message $GREEN "✓ 系统基础数据导出完成"
}

# 导出业务基础数据
export_business_data() {
    local output_file="$EXPORT_DIR/02_business_data.sql"
    local business_tables=(
        "sys_vehicle_type" "airport_areas" "spatial_rules"
        "spatial_rule_vehicle_types" "transport_routes"
    )
    
    local table_args=""
    for table in "${business_tables[@]}"; do
        table_args="$table_args --table=$table"
    done
    
    cat > "$output_file" << 'EOF'
-- QAUP 业务基础数据
SET session_replication_role = replica;
EOF
    
    pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        --data-only --no-owner --no-privileges --disable-triggers --inserts \
        $table_args >> "$output_file" 2>/dev/null || true
    
    cat >> "$output_file" << 'EOF'
SET session_replication_role = DEFAULT;
EOF
    
    print_message $GREEN "✓ 业务基础数据导出完成"
}

# 导出示例数据
export_sample_data() {
    local output_file="$EXPORT_DIR/03_sample_data.sql"
    local sample_tables=("sys_vehicle_info")
    
    local table_args=""
    for table in "${sample_tables[@]}"; do
        table_args="$table_args --table=$table"
    done
    
    cat > "$output_file" << 'EOF'
-- QAUP 示例数据
SET session_replication_role = replica;
EOF
    
    pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
        --data-only --no-owner --no-privileges --disable-triggers --inserts \
        $table_args >> "$output_file" 2>/dev/null || true
    
    cat >> "$output_file" << 'EOF'
SET session_replication_role = DEFAULT;
EOF
    
    print_message $GREEN "✓ 示例数据导出完成"
}

# 生成合并的完整数据文件
generate_combined_data() {
    local combined_file="$EXPORT_DIR/initial_data_complete.sql"
    
    cat > "$combined_file" << 'EOF'
-- QAUP 系统完整初始数据
-- 此文件由 db-manager.sh 自动生成

EOF
    
    for data_file in "01_system_data.sql" "02_business_data.sql" "03_sample_data.sql"; do
        if [ -f "$EXPORT_DIR/$data_file" ]; then
            cat "$EXPORT_DIR/$data_file" >> "$combined_file"
            echo "" >> "$combined_file"
        fi
    done
    
    print_message $GREEN "✓ 完整初始数据文件生成完成: $combined_file"
}

# ==========================================
# 备份恢复功能
# ==========================================

# 备份数据库
backup_database() {
    print_message $BLUE "备份数据库..."
    
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local backup_dir="$PROJECT_ROOT/backup/postgres"
    local backup_file="$backup_dir/qaup_backup_$timestamp.sql"
    
    mkdir -p "$backup_dir"
    
    # 检查是否在容器内
    if [ -f "/.dockerenv" ]; then
        # 在容器内执行备份
        pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DB" \
            --no-owner --no-privileges > "$backup_file"
    else
        # 在容器外执行备份
        test_database_connection
        pg_dump -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
            --no-owner --no-privileges > "$backup_file"
    fi
    
    print_message $GREEN "✓ 数据库备份完成: $backup_file"
}

# 恢复数据库
restore_database() {
    local backup_file=$1
    
    if [ -z "$backup_file" ]; then
        print_message $RED "错误: 请指定备份文件"
        exit 1
    fi
    
    if [ ! -f "$backup_file" ]; then
        print_message $RED "错误: 备份文件不存在: $backup_file"
        exit 1
    fi
    
    print_message $BLUE "恢复数据库..."
    print_message $YELLOW "警告: 这将覆盖现有数据"
    
    read -p "确定要继续吗? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
    
    # 检查是否在容器内
    if [ -f "/.dockerenv" ]; then
        # 在容器内执行恢复
        psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -f "$backup_file"
    else
        # 在容器外执行恢复
        test_database_connection
        psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$backup_file"
    fi
    
    print_message $GREEN "✓ 数据库恢复完成"
}

# ==========================================
# 健康检查功能
# ==========================================

# 健康检查
health_check() {
    print_message $BLUE "数据库健康检查..."
    
    # 检查是否在容器内
    if [ -f "/.dockerenv" ]; then
        # 容器内健康检查
        container_health_check
    else
        # 容器外健康检查
        external_health_check
    fi
}

# 容器内健康检查
container_health_check() {
    local errors=0
    
    # 检查 PostgreSQL 进程
    if pgrep -x postgres > /dev/null; then
        print_message $GREEN "✓ PostgreSQL 进程运行正常"
    else
        print_message $RED "✗ PostgreSQL 进程未运行"
        errors=$((errors + 1))
    fi
    
    # 检查数据库连接
    if psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c "SELECT 1;" > /dev/null 2>&1; then
        print_message $GREEN "✓ 数据库连接正常"
    else
        print_message $RED "✗ 数据库连接失败"
        errors=$((errors + 1))
    fi
    
    # 检查表数量
    local table_count=$(psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -t -c "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';" 2>/dev/null | tr -d ' ')
    if [ "$table_count" -gt 0 ]; then
        print_message $GREEN "✓ 数据库包含 $table_count 个表"
    else
        print_message $RED "✗ 数据库中没有表"
        errors=$((errors + 1))
    fi
    
    if [ $errors -eq 0 ]; then
        print_message $GREEN "数据库健康检查通过"
    else
        print_message $RED "数据库健康检查失败，发现 $errors 个问题"
        exit 1
    fi
}

# 容器外健康检查
external_health_check() {
    test_database_connection
    
    # 检查表数量
    local table_count=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -c "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';" 2>/dev/null | tr -d ' ')
    if [ "$table_count" -gt 0 ]; then
        print_message $GREEN "✓ 数据库包含 $table_count 个表"
    else
        print_message $RED "✗ 数据库中没有表"
        exit 1
    fi
    
    print_message $GREEN "数据库健康检查通过"
}

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

main() {
    if [ $# -eq 0 ]; then
        show_help
        exit 0
    fi
    
    local command=$1
    shift
    
    case $command in
        # 数据准备命令
        prepare-default)
            prepare_default_data
            ;;
        prepare-export)
            prepare_export_data
            ;;
        prepare-custom)
            prepare_custom_data
            ;;
        prepare-status)
            show_data_status
            ;;
        prepare-clean)
            clean_data_files
            ;;
        
        # 导出命令
        export-all)
            export_all_data
            ;;
        export-schema)
            export_schema_only
            ;;
        export-initial)
            export_initial_data
            ;;
        
        # 备份恢复命令
        backup)
            backup_database
            ;;
        restore)
            restore_database "$@"
            ;;
        
        # 维护命令
        health)
            health_check
            ;;
        test-connection)
            test_database_connection
            ;;
        
        # 帮助
        help|--help|-h)
            show_help
            ;;
        
        *)
            print_message $RED "未知命令: $command"
            show_help
            exit 1
            ;;
    esac
}

main "$@"