#!/bin/bash

# 压力测试快速启动脚本

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

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

function log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

function log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

function log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

function show_usage() {
    cat << EOF
压力测试脚本（优化版）

用法:
    $0 [command] [options]

命令:
    data        生成测试数据（使用优化的批量生成）
    test        执行压力测试（支持真实认证）
    monitor     监控系统性能
    clean       清理测试数据
    report      生成压测报告
    all         执行完整流程（生成数据 + 压测 + 报告）

选项:
    数据生成:
      -u, --users NUM         生成用户数量（默认: 1000）
      -l, --listings NUM      生成商品数量（默认: 5000）
      -o, --orders NUM        生成订单数量（默认: 3000）
      --ledger NUM            生成钱包流水数量（默认: 10000）
      --chat-messages NUM     生成聊天消息数量（默认: 5000）
      --use-optimized         使用优化的数据生成脚本（推荐）

    压力测试:
      -c, --concurrency NUM   并发数（默认: 50）
      -d, --duration SEC      测试时长/秒（默认: 60）
      -s, --scenario NAME     测试场景:
                                - realistic: 真实业务场景（默认）
                                - admin: 管理后台场景
                                - listing_only: 仅商品查询
      --gradual               启用梯度压测
      --warmup NUM            预热用户数（默认: 100）
      --use-improved          使用改进的压测工具（支持真实认证）
      --url URL               后端地址（默认: http://localhost:8080）

    通用:
      -h, --help              显示帮助信息

示例:
    # 小规模数据生成（使用优化脚本）
    $0 data --use-optimized

    # 真实场景压测（100个用户token，50并发，持续60秒）
    $0 test --use-improved -c 50 -d 60 -s realistic --warmup 100

    # 管理后台压测
    $0 test --use-improved -c 30 -d 120 -s admin

    # 梯度压测（逐步增加负载）
    $0 test --use-improved --gradual -d 60

    # 执行完整流程（优化版）
    $0 all --use-optimized --use-improved

    # 清理测试数据
    $0 clean

EOF
}

# 默认参数
USERS=1000
LISTINGS=5000
ORDERS=3000
LEDGER=10000
CHAT_MESSAGES=5000
CONCURRENCY=50
DURATION=60
SCENARIO="realistic"
BASE_URL="http://localhost:8080"
USE_OPTIMIZED=false
USE_IMPROVED=true
GRADUAL=false
WARMUP=100

# 解析命令行参数
COMMAND=""
while [[ $# -gt 0 ]]; do
    case $1 in
        data|test|monitor|clean|report|all)
            COMMAND="$1"
            shift
            ;;
        -u|--users)
            USERS="$2"
            shift 2
            ;;
        -l|--listings)
            LISTINGS="$2"
            shift 2
            ;;
        -o|--orders)
            ORDERS="$2"
            shift 2
            ;;
        --ledger)
            LEDGER="$2"
            shift 2
            ;;
        --chat-messages)
            CHAT_MESSAGES="$2"
            shift 2
            ;;
        -c|--concurrency)
            CONCURRENCY="$2"
            shift 2
            ;;
        -d|--duration)
            DURATION="$2"
            shift 2
            ;;
        -s|--scenario)
            SCENARIO="$2"
            shift 2
            ;;
        --url)
            BASE_URL="$2"
            shift 2
            ;;
        --use-optimized)
            USE_OPTIMIZED=true
            shift
            ;;
        --use-improved)
            USE_IMPROVED=true
            shift
            ;;
        --gradual)
            GRADUAL=true
            shift
            ;;
        --warmup)
            WARMUP="$2"
            shift 2
            ;;
        -h|--help)
            show_usage
            exit 0
            ;;
        *)
            log_error "未知参数: $1"
            show_usage
            exit 1
            ;;
    esac
done

if [ -z "$COMMAND" ]; then
    log_error "请指定命令"
    show_usage
    exit 1
fi

# 检查数据库连接
function check_database() {
    log_info "检查数据库连接..."
    if docker exec hfb-mysql mysql -uhfb -psecret -e "SELECT 1" &>/dev/null; then
        log_info "数据库连接正常"
        return 0
    else
        log_error "数据库连接失败，请确保 Docker 容器正在运行"
        log_info "提示: 执行 'docker-compose -f deploy/docker-compose.dev.yml up -d' 启动服务"
        return 1
    fi
}

# 检查后端服务
function check_backend() {
    log_info "检查后端服务..."
    if curl -s "$BASE_URL/health" &>/dev/null; then
        log_info "后端服务正常"
        return 0
    else
        log_warn "后端服务未响应: $BASE_URL"
        log_info "提示: 确保后端服务已启动"
        return 1
    fi
}

# 生成测试数据
function generate_data() {
    log_info "开始生成测试数据..."
    log_info "配置: 用户=$USERS, 商品=$LISTINGS, 订单=$ORDERS, 钱包流水=$LEDGER, 聊天消息=$CHAT_MESSAGES"

    check_database || exit 1

    # 选择使用的SQL脚本
    SQL_FILE="$SCRIPT_DIR/load_test_data.sql"
    if [ "$USE_OPTIMIZED" = true ]; then
        SQL_FILE="$SCRIPT_DIR/load_test_data_optimized.sql"
        log_info "使用优化的数据生成脚本"
    fi

    # 创建临时SQL文件
    TMP_SQL="/tmp/load_test_data_$(date +%s).sql"
    trap 'rm -f "$TMP_SQL"' RETURN

    if [ "$USE_OPTIMIZED" = true ]; then
        cat > "$TMP_SQL" << EOF
-- 优化版数据生成
USE hfb_sys;

-- 调用批量生成存储过程
CALL generate_users_batch($USERS);
CALL generate_listings_batch($LISTINGS);
CALL generate_orders_batch($ORDERS);
CALL generate_wallet_ledger_batch($LEDGER);
CALL generate_chat_data_batch($CHAT_MESSAGES);

-- 显示统计
SELECT '用户数' as item, COUNT(*) as count FROM users
UNION ALL SELECT '游戏账号数', COUNT(*) FROM game_accounts
UNION ALL SELECT '商品数', COUNT(*) FROM rental_listings
UNION ALL SELECT '订单数', COUNT(*) FROM rental_orders
UNION ALL SELECT '钱包流水数', COUNT(*) FROM wallet_ledger
UNION ALL SELECT '聊天会话数', COUNT(*) FROM chat_conversations
UNION ALL SELECT '聊天消息数', COUNT(*) FROM chat_messages;
EOF
    else
        cat > "$TMP_SQL" << EOF
-- 原始版数据生成
USE hfb_sys;

-- 调用存储过程生成数据
CALL generate_users($USERS);
CALL generate_listings($LISTINGS);
CALL generate_orders($ORDERS);
CALL generate_wallet_ledger($LEDGER);

-- 生成聊天会话
INSERT INTO chat_conversations (order_id, type, title, status, last_message_at)
SELECT id, 'order_group', CONCAT('订单', order_no, '群聊'), 'active', created_at
FROM rental_orders
WHERE id <= 1000
ON DUPLICATE KEY UPDATE order_id=order_id;

CALL generate_chat_messages($CHAT_MESSAGES);

-- 显示统计
SELECT '用户数' as item, COUNT(*) as count FROM users
UNION ALL SELECT '游戏账号数', COUNT(*) FROM game_accounts
UNION ALL SELECT '商品数', COUNT(*) FROM rental_listings
UNION ALL SELECT '订单数', COUNT(*) FROM rental_orders
UNION ALL SELECT '钱包流水数', COUNT(*) FROM wallet_ledger
UNION ALL SELECT '聊天会话数', COUNT(*) FROM chat_conversations
UNION ALL SELECT '聊天消息数', COUNT(*) FROM chat_messages;
EOF
    fi

    log_info "执行数据生成..."

    # 先执行基础SQL创建存储过程
    docker exec -i hfb-mysql mysql -uhfb -psecret hfb_sys < "$SQL_FILE"

    # 执行数据生成
    docker exec -i hfb-mysql mysql -uhfb -psecret hfb_sys < "$TMP_SQL"

    rm -f "$TMP_SQL"
    trap - RETURN

    log_info "测试数据生成完成！"
}

# 执行压力测试
function run_stress_test() {
    log_info "开始压力测试..."
    log_info "配置: 场景=$SCENARIO, 并发=$CONCURRENCY, 时长=${DURATION}秒, 目标=$BASE_URL"

    check_backend || log_warn "后端服务未响应，测试可能失败"

    # 编译压测工具
    log_info "编译压测工具..."
    cd "$SCRIPT_DIR"

    go build -o stress_test load_stress.go

    if [ $? -ne 0 ]; then
        log_error "压测工具编译失败"
        return 1
    fi

    # 执行压测
    log_info "开始执行压力测试..."

    if [ "$GRADUAL" = true ]; then
        ./stress_test -url "$BASE_URL" -c "$CONCURRENCY" -d "$DURATION" -s "$SCENARIO" -warmup "$WARMUP" -gradual
    else
        ./stress_test -url "$BASE_URL" -c "$CONCURRENCY" -d "$DURATION" -s "$SCENARIO" -warmup "$WARMUP"
    fi

    log_info "压力测试完成！"
}

# 监控系统性能
function monitor_system() {
    log_info "开始监控系统性能（按 Ctrl+C 停止）..."

    echo ""
    echo "=== Docker 容器资源使用 ==="
    docker stats --no-stream hfb-backend hfb-mysql hfb-redis 2>/dev/null || log_warn "部分容器未运行"

    echo ""
    echo "=== MySQL 连接数 ==="
    docker exec hfb-mysql mysql -uhfb -psecret -e "SHOW STATUS LIKE 'Threads_connected';" 2>/dev/null

    echo ""
    echo "=== MySQL 慢查询统计 ==="
    docker exec hfb-mysql mysql -uhfb -psecret -e "SHOW STATUS LIKE 'Slow_queries';" 2>/dev/null

    echo ""
    echo "=== Redis 信息 ==="
    docker exec hfb-redis redis-cli INFO stats | grep -E "total_connections_received|total_commands_processed|instantaneous_ops_per_sec" 2>/dev/null

    echo ""
    log_info "持续监控请使用: docker stats"
}

# 清理测试数据
function clean_data() {
    log_warn "即将清理所有测试数据，此操作不可恢复！"
    read -p "确认继续？(yes/no): " confirm

    if [ "$confirm" != "yes" ]; then
        log_info "已取消清理操作"
        exit 0
    fi

    check_database || exit 1

    log_info "开始清理测试数据..."

    docker exec -i hfb-mysql mysql -uhfb -psecret hfb_sys << 'EOF'
SET FOREIGN_KEY_CHECKS = 0;

-- 清理测试数据（保留ID < 1000的数据）
DELETE FROM chat_messages WHERE id > 100;
DELETE FROM chat_conversations WHERE id > 100;
DELETE FROM wallet_ledger WHERE id > 100;
DELETE FROM rental_orders WHERE id > 100;
DELETE FROM rental_listings WHERE id > 100;
DELETE FROM game_accounts WHERE id > 100;
DELETE FROM users WHERE id > 1000;

-- 重置自增ID
ALTER TABLE users AUTO_INCREMENT = 1001;
ALTER TABLE game_accounts AUTO_INCREMENT = 101;
ALTER TABLE rental_listings AUTO_INCREMENT = 101;
ALTER TABLE rental_orders AUTO_INCREMENT = 101;
ALTER TABLE wallet_ledger AUTO_INCREMENT = 101;
ALTER TABLE chat_conversations AUTO_INCREMENT = 101;
ALTER TABLE chat_messages AUTO_INCREMENT = 101;

SET FOREIGN_KEY_CHECKS = 1;

-- 显示清理后的统计
SELECT '用户数' as item, COUNT(*) as count FROM users
UNION ALL SELECT '商品数', COUNT(*) FROM rental_listings
UNION ALL SELECT '订单数', COUNT(*) FROM rental_orders
UNION ALL SELECT '钱包流水数', COUNT(*) FROM wallet_ledger;
EOF

    log_info "测试数据清理完成！"
}

# 生成压测报告
function generate_report() {
    log_info "生成压测报告..."

    REPORT_FILE="$PROJECT_ROOT/stress_test_report_$(date +%Y%m%d_%H%M%S).txt"

    {
        echo "==================================="
        echo "压力测试报告"
        echo "生成时间: $(date '+%Y-%m-%d %H:%M:%S')"
        echo "==================================="
        echo ""

        echo "--- 数据库统计 ---"
        docker exec hfb-mysql mysql -uhfb -psecret hfb_sys -e "
            SELECT '用户数' as item, COUNT(*) as count FROM users
            UNION ALL SELECT '商品数', COUNT(*) FROM rental_listings
            UNION ALL SELECT '订单数', COUNT(*) FROM rental_orders
            UNION ALL SELECT '钱包流水', COUNT(*) FROM wallet_ledger;
        " 2>/dev/null

        echo ""
        echo "--- MySQL 性能指标 ---"
        docker exec hfb-mysql mysql -uhfb -psecret -e "
            SHOW STATUS LIKE 'Threads_connected';
            SHOW STATUS LIKE 'Slow_queries';
            SHOW STATUS LIKE 'Questions';
        " 2>/dev/null

        echo ""
        echo "--- 表大小统计 ---"
        docker exec hfb-mysql mysql -uhfb -psecret hfb_sys -e "
            SELECT
                table_name AS '表名',
                ROUND(((data_length + index_length) / 1024 / 1024), 2) AS '大小(MB)',
                table_rows AS '行数'
            FROM information_schema.TABLES
            WHERE table_schema = 'hfb_sys'
            ORDER BY (data_length + index_length) DESC
            LIMIT 10;
        " 2>/dev/null

    } > "$REPORT_FILE"

    cat "$REPORT_FILE"

    log_info "报告已保存到: $REPORT_FILE"
}

# 执行完整流程
function run_all() {
    log_info "执行完整压测流程..."

    generate_data
    sleep 3

    run_stress_test
    sleep 2

    generate_report

    log_info "完整流程执行完成！"
}

# 执行命令
case $COMMAND in
    data)
        generate_data
        ;;
    test)
        run_stress_test
        ;;
    monitor)
        monitor_system
        ;;
    clean)
        clean_data
        ;;
    report)
        generate_report
        ;;
    all)
        run_all
        ;;
    *)
        log_error "未知命令: $COMMAND"
        show_usage
        exit 1
        ;;
esac