#!/bin/bash
#
# 水务管理系统 - 统一性能压力测试入口脚本
#
# 依次运行：REST API / WebSocket / MQTT IoT / 数据库查询 四项压力测试，
# 收集系统指标（CPU、内存、磁盘 IO），生成综合报告。
#
# 用法：
#   bash run_all_benchmarks.sh [--quick] [--output-dir DIR]
#
# --quick: 快速模式（减少测试时长，适合 CI/CD）
# --output-dir: 结果输出目录（默认 ./reports/YYYYMMDD_HHMMSS）
#

set -euo pipefail

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

# ==================== 参数解析 ====================
QUICK_MODE=false
OUTPUT_DIR=""
API_HOST="${API_HOST:-localhost}"
API_PORT="${API_PORT:-8080}"
WS_HOST="${WS_HOST:-localhost}"
WS_PORT="${WS_PORT:-8765}"
MQTT_BROKER="${MQTT_BROKER:-localhost}"
MQTT_PORT="${MQTT_PORT:-1883}"
DB_HOST="${DB_HOST:-localhost}"
DB_PORT="${DB_PORT:-5432}"
DB_NAME="${DB_NAME:-water_management}"
DB_USER="${DB_USER:-postgres}"
DB_PASS="${DB_PASS:-postgres}"

while [[ $# -gt 0 ]]; do
    case $1 in
        --quick) QUICK_MODE=true; shift ;;
        --output-dir) OUTPUT_DIR="$2"; shift 2 ;;
        --api-host) API_HOST="$2"; shift 2 ;;
        --ws-host) WS_HOST="$2"; shift 2 ;;
        --mqtt-broker) MQTT_BROKER="$2"; shift 2 ;;
        --db-host) DB_HOST="$2"; shift 2 ;;
        *) echo "未知参数: $1"; exit 1 ;;
    esac
done

# 默认输出目录
if [[ -z "$OUTPUT_DIR" ]]; then
    OUTPUT_DIR="./reports/$(date +%Y%m%d_%H%M%S)"
fi
mkdir -p "$OUTPUT_DIR"

# ==================== 工具函数 ====================

log() { echo "[$(date '+%H:%M:%S')] $*"; }
error() { echo "[$(date '+%H:%M:%S')] ❌ $*" >&2; }

check_python_deps() {
    log "检查 Python 依赖..."
    pip install -q -r requirements.txt 2>/dev/null || {
        error "依赖安装失败，请手动执行: pip install -r requirements.txt"
        exit 1
    }
}

# 收集系统指标
SYSMONITOR_PID=""
start_system_monitor() {
    local interval=5
    local output_file="$1"
    (
        echo "timestamp,cpu_percent,memory_mb,disk_read_mb,disk_write_mb" > "$output_file"
        while true; do
            local ts=$(date '+%Y-%m-%d %H:%M:%S')
            local cpu=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' 2>/dev/null || echo "0")
            local mem=$(free -m | awk '/Mem:/{print $3}' 2>/dev/null || echo "0")
            local dr=$(cat /proc/diskstats 2>/dev/null | awk '{r+=$6} END{print r/2048}' || echo "0")
            local dw=$(cat /proc/diskstats 2>/dev/null | awk '{w+=$10} END{print w/2048}' || echo "0")
            echo "$ts,$cpu,$mem,$dr,$dw" >> "$output_file"
            sleep $interval
        done
    ) &
    SYSMONITOR_PID=$!
}

stop_system_monitor() {
    if [[ -n "$SYSMONITOR_PID" ]]; then
        kill "$SYSMONITOR_PID" 2>/dev/null || true
        wait "$SYSMONITOR_PID" 2>/dev/null || true
        SYSMONITOR_PID=""
    fi
}

# 收集系统信息
collect_system_info() {
    local output_file="$1"
    cat > "$output_file" << EOF
# 系统环境信息
- 主机名: $(hostname)
- 操作系统: $(uname -a)
- CPU: $(nproc) 核
- 内存: $(free -h | awk '/Mem:/{print $2}')
- 磁盘: $(df -h / | awk 'NR==2{print $2 " total, " $4 " available"}')
- Python: $(python3 --version 2>&1)
- 测试时间: $(date '+%Y-%m-%d %H:%M:%S')
- 快速模式: $QUICK_MODE
EOF
}

# ==================== 测试配置 ====================

if $QUICK_MODE; then
    REST_USERS=50
    REST_DURATION="30s"
    WS_CLIENTS=100
    WS_DURATION=30
    MQTT_DEVICES=100
    MQTT_DURATION=30
    DB_RECORDS=100000
else
    REST_USERS=500
    REST_DURATION="3m"
    WS_CLIENTS=1000
    WS_DURATION=120
    MQTT_DEVICES=1000
    MQTT_DURATION=120
    DB_RECORDS=1000000
fi

# ==================== 开始测试 ====================

echo ""
echo "============================================================"
echo "🏗️  水务管理系统 - 性能压力测试"
echo "============================================================"
echo "测试时间: $(date '+%Y-%m-%d %H:%M:%S')"
echo "输出目录: $OUTPUT_DIR"
echo "快速模式: $QUICK_MODE"
echo "============================================================"

# 检查依赖
check_python_deps

# 收集系统信息
SYSTEM_INFO="$OUTPUT_DIR/system_info.md"
collect_system_info "$SYSTEM_INFO"

# 启动系统监控
SYSMON_CSV="$OUTPUT_DIR/system_metrics.csv"
start_system_monitor "$SYSMON_CSV"

# 注册退出清理
trap stop_system_monitor EXIT

# ---------- 1. REST API 压力测试 ----------
log ""
log "=========================================="
log "📡 测试 1/4: REST API 压力测试"
log "=========================================="

REST_RESULT="$OUTPUT_DIR/rest_api_result.json"
if command -v locust &>/dev/null; then
    timeout 600 locust -f locustfile.py \
        --host="http://${API_HOST}:${API_PORT}" \
        --headless -u "$REST_USERS" -r 10 \
        --run-time "$REST_DURATION" \
        --csv="$OUTPUT_DIR/rest_api" \
        --json 2>&1 | tee "$OUTPUT_DIR/rest_api.log" || true
    log "✅ REST API 测试完成"
else
    log "⚠️ locust 未安装，跳过 REST API 测试"
    echo '{"skipped": true, "reason": "locust not installed"}' > "$REST_RESULT"
fi

# ---------- 2. WebSocket 压力测试 ----------
log ""
log "=========================================="
log "🔌 测试 2/4: WebSocket 压力测试"
log "=========================================="

WS_RESULT="$OUTPUT_DIR/websocket_result.json"
timeout 600 python3 websocket_stress.py \
    --host "$WS_HOST" --port "$WS_PORT" \
    --clients "$WS_CLIENTS" --duration "$WS_DURATION" \
    --output "$WS_RESULT" 2>&1 | tee "$OUTPUT_DIR/websocket.log" || true
log "✅ WebSocket 测试完成"

# ---------- 3. MQTT IoT 压力测试 ----------
log ""
log "=========================================="
log "📡 测试 3/4: IoT MQTT 压力测试"
log "=========================================="

MQTT_RESULT="$OUTPUT_DIR/mqtt_result.json"
timeout 600 python3 mqtt_iot_stress.py \
    --broker "$MQTT_BROKER" --port "$MQTT_PORT" \
    --devices "$MQTT_DEVICES" --duration "$MQTT_DURATION" \
    --output "$MQTT_RESULT" 2>&1 | tee "$OUTPUT_DIR/mqtt.log" || true
log "✅ MQTT 测试完成"

# ---------- 4. 数据库查询测试 ----------
log ""
log "=========================================="
log "🗄️ 测试 4/4: 数据库查询压力测试"
log "=========================================="

DB_RESULT="$OUTPUT_DIR/db_query_result.json"
timeout 600 python3 db_query_stress.py \
    --host "$DB_HOST" --port "$DB_PORT" \
    --db "$DB_NAME" --user "$DB_USER" --password "$DB_PASS" \
    --generate-data "$DB_RECORDS" \
    --index-compare \
    --output "$DB_RESULT" 2>&1 | tee "$OUTPUT_DIR/db_query.log" || true
log "✅ 数据库测试完成"

# ==================== 停止监控 ====================
stop_system_monitor

# ==================== 生成综合报告 ====================
log ""
log "=========================================="
log "📝 生成综合报告"
log "=========================================="

REPORT_FILE="$OUTPUT_DIR/report.md"

# 复制报告模板并填充
if [[ -f report_template.md ]]; then
    cp report_template.md "$REPORT_FILE"
    # 追加实际结果
    cat >> "$REPORT_FILE" << EOF

---
## 实际测试结果

### 测试环境
$(cat "$SYSTEM_INFO")

### REST API 结果
\`\`\`
$(tail -30 "$OUTPUT_DIR/rest_api.log" 2>/dev/null || echo "无数据")
\`\`\`

### WebSocket 结果
\`\`\`json
$(cat "$WS_RESULT" 2>/dev/null || echo "无数据")
\`\`\`

### MQTT IoT 结果
\`\`\`json
$(cat "$MQTT_RESULT" 2>/dev/null || echo "无数据")
\`\`\`

### 数据库查询结果
\`\`\`json
$(cat "$DB_RESULT" 2>/dev/null || echo "无数据")
\`\`\`

### 系统资源指标
\`\`\`csv
$(head -20 "$SYSMON_CSV" 2>/dev/null || echo "无数据")
\`\`\`
EOF
else
    # 简单报告
    cat > "$REPORT_FILE" << EOF
# 性能压力测试报告

## 测试时间
$(date '+%Y-%m-%d %H:%M:%S')

## 测试环境
$(cat "$SYSTEM_INFO")

## 测试结果摘要

| 测试项 | 状态 | 结果文件 |
|--------|------|----------|
| REST API | ✅ | rest_api_result.json |
| WebSocket | ✅ | websocket_result.json |
| MQTT IoT | ✅ | mqtt_result.json |
| 数据库查询 | ✅ | db_query_result.json |

## 系统资源指标
见 system_metrics.csv
EOF
fi

log "✅ 综合报告已生成: $REPORT_FILE"

# ==================== 完成 ====================
echo ""
echo "============================================================"
echo "🎉 所有压力测试完成！"
echo "============================================================"
echo "输出目录: $OUTPUT_DIR"
echo "文件列表:"
ls -la "$OUTPUT_DIR/"
echo "============================================================"