#!/bin/bash

# 达梦8数据库性能测试脚本
# 测试目标：
# 1. 并发连接数 ≥ 5000个
# 2. SQL语句解析能力 ≥ 3000条/秒


# 脚本说明：
# 小规模测试模式
# ./dm8-test.sh -t

# 自定义测试
# ./dm8-test.sh -c 1000 -d 30

# 完整测试（5000连接）
# ./dm8-test.sh

# 配置参数
API_URL="http://127.0.0.1:55000/api/DMHR.EMPLOYEE"
MAX_CONCURRENT_CONNECTIONS=5000
TARGET_SQL_RATE=3000 # 每秒SQL执行数
TEST_DURATION=60     # 测试持续时间（秒）
TEMP_DIR="/tmp/dm8_test"
RESULTS_FILE="$TEMP_DIR/test_results.log"

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

# 创建临时目录
mkdir -p "$TEMP_DIR"
>"$RESULTS_FILE"

# 日志函数
log() {
    echo -e "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$RESULTS_FILE"
}

# 错误处理
error_exit() {
    log "${RED}错误: $1${NC}"
    cleanup
    exit 1
}

# 清理函数
cleanup() {
    log "${YELLOW}正在清理测试环境...${NC}"
    # 杀死所有后台进程
    jobs -p | xargs -r kill -9 2>/dev/null
    # 清理临时文件
    rm -f "$TEMP_DIR"/worker_*.log 2>/dev/null
    log "${GREEN}清理完成${NC}"
}

# 信号处理
trap cleanup EXIT INT TERM

# 检查依赖
check_dependencies() {
    log "${BLUE}检查依赖工具...${NC}"

    if ! command -v curl &>/dev/null; then
        error_exit "curl 未安装，请先安装 curl"
    fi

    if ! command -v bc &>/dev/null; then
        error_exit "bc 未安装，请先安装 bc"
    fi

    log "${GREEN}依赖检查通过${NC}"
}

# 测试API连通性
test_connectivity() {
    log "${BLUE}测试API连通性...${NC}"

    local response_code=$(curl -s -w "%{http_code}" --connect-timeout 5 --max-time 10 "$API_URL" -o /dev/null 2>/dev/null)

    if [ "$response_code" = "200" ]; then
        log "${GREEN}API连通性测试通过 (HTTP $response_code)${NC}"
    elif [ "$response_code" = "000" ]; then
        log "${YELLOW}警告: 无法获取HTTP响应码，但连接可能成功${NC}"
    else
        log "${YELLOW}警告: API返回HTTP $response_code，但将继续测试${NC}"
    fi

    # 进行小规模预测试
    log "${BLUE}进行小规模预测试（10个连接）...${NC}"
    local test_success=0
    for i in {1..10}; do
        if curl -s --connect-timeout 2 --max-time 5 "$API_URL" >/dev/null 2>&1; then
            ((test_success++))
        fi
    done

    log "预测试结果: $test_success/10 成功"
    if [ "$test_success" -eq 0 ]; then
        log "${YELLOW}警告: 预测试全部失败，但将继续进行大规模测试${NC}"
    fi
}

# 单个工作进程
worker_process() {
    local worker_id=$1
    local worker_log="$TEMP_DIR/worker_${worker_id}.log"
    local request_count=0
    local success_count=0
    local error_count=0
    local start_time=$(date +%s)
    local end_time=$((start_time + TEST_DURATION))

    while [ $(date +%s) -lt $end_time ]; do
        local response_code=$(curl -s -w "%{http_code}" --connect-timeout 2 --max-time 5 "$API_URL" -o /dev/null 2>/dev/null)

        if [ "$response_code" = "200" ] || [ "$response_code" = "000" ]; then
            ((success_count++))
        else
            ((error_count++))
        fi
        ((request_count++))

        # 控制请求频率，避免过快
        sleep 0.01
    done

    echo "$worker_id:$request_count:$success_count:$error_count" >"$worker_log"
}

# 并发连接测试
test_concurrent_connections() {
    log "${BLUE}开始并发连接测试...${NC}"
    log "目标并发连接数: $MAX_CONCURRENT_CONNECTIONS"
    log "测试持续时间: ${TEST_DURATION}秒"

    local start_time=$(date +%s)

    # 启动工作进程
    for ((i = 1; i <= MAX_CONCURRENT_CONNECTIONS; i++)); do
        worker_process $i &

        # 每启动100个进程显示一次进度
        if ((i % 100 == 0)); then
            echo -ne "\r启动进程: $i/$MAX_CONCURRENT_CONNECTIONS"
        fi
    done
    echo ""

    log "${YELLOW}已启动 $MAX_CONCURRENT_CONNECTIONS 个并发连接${NC}"
    log "${YELLOW}等待测试完成（${TEST_DURATION}秒）...${NC}"

    # 等待所有工作进程完成
    wait

    local end_time=$(date +%s)
    local actual_duration=$((end_time - start_time))

    log "${GREEN}并发连接测试完成，实际耗时: ${actual_duration}秒${NC}"
}

# 计算测试结果
calculate_results() {
    log "${BLUE}计算测试结果...${NC}"

    local total_requests=0
    local total_success=0
    local active_connections=0

    # 统计所有工作进程的结果
    for log_file in "$TEMP_DIR"/worker_*.log; do
        if [ -f "$log_file" ]; then
            local line=$(cat "$log_file")
            local worker_id=$(echo "$line" | cut -d: -f1)
            local requests=$(echo "$line" | cut -d: -f2)
            local success=$(echo "$line" | cut -d: -f3)
            local errors=0

            # 检查是否有错误计数字段
            if [ "$(echo "$line" | grep -o ':' | wc -l)" -eq 3 ]; then
                errors=$(echo "$line" | cut -d: -f4)
            fi

            # 验证数字
            if [[ "$requests" =~ ^[0-9]+$ ]] && [[ "$success" =~ ^[0-9]+$ ]]; then
                total_requests=$((total_requests + requests))
                total_success=$((total_success + success))

                # 如果有成功请求，说明连接建立成功
                if [ "$success" -gt 0 ]; then
                    ((active_connections++))
                fi
            fi
        fi
    done

    # 计算性能指标
    local success_rate=0
    local avg_sql_rate=0

    if [ "$total_requests" -gt 0 ]; then
        success_rate=$(echo "scale=2; $total_success * 100 / $total_requests" | bc)
    fi

    if [ "$TEST_DURATION" -gt 0 ]; then
        avg_sql_rate=$(echo "scale=2; $total_success / $TEST_DURATION" | bc)
    fi

    # 输出结果
    log "${GREEN}===== 测试结果 =====${NC}"
    log "总请求数: $total_requests"
    log "成功请求数: $total_success"
    log "成功率: ${success_rate}%"
    log "有效并发连接数: $active_connections"
    log "平均SQL执行率: ${avg_sql_rate} 条/秒"
    log ""

    # 评估测试结果
    log "${BLUE}===== 性能评估 =====${NC}"

    # 并发连接数评估
    if [ "$active_connections" -ge 5000 ]; then
        log "${GREEN}✓ 并发连接数测试通过: $active_connections ≥ 5000${NC}"
    else
        log "${RED}✗ 并发连接数测试失败: $active_connections < 5000${NC}"
    fi

    # SQL解析能力评估
    local sql_rate_int=0
    if [ -n "$avg_sql_rate" ] && [[ "$avg_sql_rate" =~ ^[0-9]+\.?[0-9]*$ ]]; then
        sql_rate_int=$(echo "$avg_sql_rate" | cut -d. -f1)
        # 如果没有整数部分，设为0
        if [ -z "$sql_rate_int" ]; then
            sql_rate_int=0
        fi
    fi

    if [ "$sql_rate_int" -ge 3000 ]; then
        log "${GREEN}✓ SQL解析能力测试通过: ${avg_sql_rate} ≥ 3000 条/秒${NC}"
    else
        log "${RED}✗ SQL解析能力测试失败: ${avg_sql_rate} < 3000 条/秒${NC}"
    fi

    # 整体评估
    if [ "$active_connections" -ge 5000 ] && [ "$sql_rate_int" -ge 3000 ]; then
        log "${GREEN}🎉 达梦8数据库性能测试全部通过！${NC}"
        return 0
    else
        log "${RED}❌ 达梦8数据库性能测试未完全通过${NC}"
        return 1
    fi
}

# 主函数
main() {
    log "${GREEN}===== 达梦8数据库性能测试开始 =====${NC}"
    log "测试目标："
    log "  1. 并发连接数 ≥ 5000个"
    log "  2. SQL语句解析能力 ≥ 3000条/秒"
    log ""

    # 检查系统资源
    log "${BLUE}系统信息:${NC}"
    log "CPU核心数: $(nproc)"
    log "内存大小: $(free -h | grep '^Mem:' | awk '{print $2}')"
    log "当前时间: $(date)"
    log ""

    # 执行测试步骤
    check_dependencies
    test_connectivity

    # 调整系统参数以支持大量并发连接
    log "${YELLOW}调整系统参数以支持高并发...${NC}"
    ulimit -n 65536 2>/dev/null || log "${YELLOW}警告: 无法调整文件描述符限制${NC}"

    test_concurrent_connections
    calculate_results

    local exit_code=$?

    log "${GREEN}===== 测试完成 =====${NC}"
    log "详细日志保存在: $RESULTS_FILE"

    return $exit_code
}

# 显示帮助信息
show_help() {
    echo "达梦8数据库性能测试脚本"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  -h, --help              显示此帮助信息"
    echo "  -c, --connections NUM   设置最大并发连接数 (默认: 5000)"
    echo "  -d, --duration SEC      设置测试持续时间 (默认: 60秒)"
    echo "  -u, --url URL          设置API URL (默认: $API_URL)"
    echo "  -t, --test-mode         小规模测试模式 (100个连接，10秒)"
    echo ""
    echo "示例:"
    echo "  $0                      使用默认参数运行测试"
    echo "  $0 -c 6000 -d 120      测试6000个并发连接，持续120秒"
    echo "  $0 -t                   运行小规模测试模式"
    echo ""
}

# 参数解析
while [[ $# -gt 0 ]]; do
    case $1 in
    -h | --help)
        show_help
        exit 0
        ;;
    -c | --connections)
        MAX_CONCURRENT_CONNECTIONS="$2"
        shift 2
        ;;
    -d | --duration)
        TEST_DURATION="$2"
        shift 2
        ;;
    -u | --url)
        API_URL="$2"
        shift 2
        ;;
    -t | --test-mode)
        MAX_CONCURRENT_CONNECTIONS=100
        TEST_DURATION=10
        log "${YELLOW}启用测试模式: 100个连接，持续10秒${NC}"
        shift
        ;;
    *)
        log "${RED}未知参数: $1${NC}"
        show_help
        exit 1
        ;;
    esac
done

# 验证参数
if ! [[ "$MAX_CONCURRENT_CONNECTIONS" =~ ^[0-9]+$ ]] || [ "$MAX_CONCURRENT_CONNECTIONS" -le 0 ]; then
    error_exit "并发连接数必须是正整数"
fi

if ! [[ "$TEST_DURATION" =~ ^[0-9]+$ ]] || [ "$TEST_DURATION" -le 0 ]; then
    error_exit "测试持续时间必须是正整数"
fi

# 运行主程序
main
exit $?