#!/bin/bash

# VPS CPU性能测试和对比脚本
# 用法: bash <(curl -sSL https://raw.githubusercontent.com/username/vpsbench/main/cpu_bench.sh)

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

# 结果文件
RESULTS_FILE="/tmp/vps_benchmark_results.txt"
COMPARE_FILE="/tmp/vps_comparison.txt"

# 安装工具
install_tools() {
    echo "安装必要工具..."
    if command -v yum >/dev/null 2>&1; then
        yum install -y sysbench bc wget curl >/dev/null 2>&1
    elif command -v apt-get >/dev/null 2>&1; then
        apt-get update >/dev/null 2>&1
        apt-get install -y sysbench bc wget curl >/dev/null 2>&1
    fi
}

# 获取系统信息
get_system_info() {
    CPU_NAME=$(grep "model name" /proc/cpuinfo | head -1 | cut -d: -f2 | sed 's/^ *//')
    CPU_CORES=$(nproc)
    MEMORY=$(free -h | grep Mem | awk '{print $2}')
    OS=$(cat /etc/os-release | grep PRETTY_NAME | cut -d= -f2 | tr -d '"')
    HOSTNAME=$(hostname)
}

# 性能测试
run_benchmarks() {
    echo "开始性能测试..."
    
    # 单核测试
    SINGLE_SCORE=$(sysbench cpu --threads=1 --cpu-max-prime=20000 run 2>/dev/null | grep "events per second" | awk '{printf "%.2f", $4}')
    
    # 多核测试
    MULTI_SCORE=$(sysbench cpu --threads=$CPU_CORES --cpu-max-prime=20000 run 2>/dev/null | grep "events per second" | awk '{printf "%.2f", $4}')
    
    # 浮点测试
    START_TIME=$(date +%s.%N)
    echo "scale=3000; 4*a(1)" | bc -l -q >/dev/null 2>&1
    END_TIME=$(date +%s.%N)
    FLOAT_TIME=$(echo "$END_TIME - $START_TIME" | bc)
    
    # 内存测试
    START_TIME=$(date +%s.%N)
    dd if=/dev/zero of=/tmp/test.bin bs=1M count=100 status=none 2>/dev/null
    END_TIME=$(date +%s.%N)
    MEM_WRITE=$(echo "scale=2; 100 / ($END_TIME - $START_TIME)" | bc)
    rm -f /tmp/test.bin
}

# 保存结果
save_results() {
    cat > $RESULTS_FILE << EOF
VPS名称: $HOSTNAME
测试时间: $(date)
CPU型号: $CPU_NAME
CPU核心: $CPU_CORES
内存大小: $MEMORY
系统版本: $OS
单核得分: $SINGLE_SCORE
多核得分: $MULTI_SCORE
浮点耗时: ${FLOAT_TIME}秒
内存写入: ${MEM_WRITE} MB/s
EOF

    echo "结果已保存到: $RESULTS_FILE"
}

# 显示当前结果
show_current_results() {
    echo ""
    echo -e "${PURPLE}=== 当前VPS测试结果 ===${NC}"
    echo "VPS名称: $HOSTNAME"
    echo "CPU: $CPU_NAME"
    echo "核心: $CPU_CORES"
    echo "单核得分: $SINGLE_SCORE"
    echo "多核得分: $MULTI_SCORE"
    echo "浮点耗时: ${FLOAT_TIME}秒"
    echo "内存写入: ${MEM_WRITE} MB/s"
}

# 性能评级
get_rating() {
    local score=$1
    local type=$2
    
    case $type in
        "single")
            if (( $(echo "$score > 800" | bc -l) )); then
                echo "${GREEN}优秀${NC}"
            elif (( $(echo "$score > 500" | bc -l) )); then
                echo "${YELLOW}良好${NC}"
            else
                echo "${RED}一般${NC}"
            fi
            ;;
        "multi")
            local expected=$(echo "$CPU_CORES * 400" | bc)
            if (( $(echo "$score > $expected" | bc -l) )); then
                echo "${GREEN}优秀${NC}"
            elif (( $(echo "$score > $(echo "$CPU_CORES * 250" | bc)" | bc -l) )); then
                echo "${YELLOW}良好${NC}"
            else
                echo "${RED}一般${NC}"
            fi
            ;;
        "float")
            if (( $(echo "$score < 5" | bc -l) )); then
                echo "${GREEN}优秀${NC}"
            elif (( $(echo "$score < 10" | bc -l) )); then
                echo "${YELLOW}良好${NC}"
            else
                echo "${RED}一般${NC}"
            fi
            ;;
    esac
}

# 显示对比参考
show_reference() {
    echo ""
    echo -e "${CYAN}=== 性能对比参考 ===${NC}"
    echo "单核性能参考:"
    echo -e "  ${GREEN}优秀: >800  ${YELLOW}良好: 500-800  ${RED}一般: <500${NC}"
    echo ""
    echo "多核性能参考 (以单核500为基准):"
    echo -e "  ${GREEN}优秀: 核心数 × 400+  ${YELLOW}良好: 核心数 × 250+  ${RED}一般: 低于核心数 × 250${NC}"
    echo ""
    echo "浮点性能参考:"
    echo -e "  ${GREEN}优秀: <5秒  ${YELLOW}良好: 5-10秒  ${RED}一般: >10秒${NC}"
}

# 添加对比数据
add_comparison() {
    echo ""
    echo -e "${YELLOW}是否将当前结果添加到对比数据库？(y/N)${NC}"
    read -r answer
    
    if [[ $answer == "y" || $answer == "Y" ]]; then
        echo "请输入这个VPS的标识名称（如: Vultr-2C4G, AWS-t2-medium）:"
        read -r vps_name
        
        if [ -n "$vps_name" ]; then
            cat >> $COMPARE_FILE << EOF
----------------------------------------
VPS名称: $vps_name
测试时间: $(date)
CPU型号: $CPU_NAME
CPU核心: $CPU_CORES
单核得分: $SINGLE_SCORE
多核得分: $MULTI_SCORE
浮点耗时: ${FLOAT_TIME}秒
内存写入: ${MEM_WRITE} MB/s
EOF
            echo -e "${GREEN}已添加到对比数据库: $COMPARE_FILE${NC}"
        fi
    fi
}

# 显示对比结果
show_comparison() {
    if [ -f "$COMPARE_FILE" ]; then
        echo ""
        echo -e "${PURPLE}=== 历史对比结果 ===${NC}"
        cat $COMPARE_FILE
    else
        echo ""
        echo -e "${YELLOW}暂无对比数据，请先添加一些测试结果${NC}"
    fi
}

# 生成对比表格
generate_comparison_table() {
    if [ ! -f "$COMPARE_FILE" ]; then
        return
    fi
    
    echo ""
    echo -e "${CYAN}=== 性能对比表格 ===${NC}"
    printf "%-20s %-8s %-12s %-12s %-12s\n" "VPS名称" "核心数" "单核得分" "多核得分" "浮点耗时"
    echo "----------------------------------------------------------------"
    
    # 从对比文件中提取数据
    while IFS= read -r line; do
        if [[ $line == "VPS名称:"* ]]; then
            vps_name=$(echo "$line" | cut -d: -f2 | sed 's/^ *//')
        elif [[ $line == "CPU核心:"* ]]; then
            cores=$(echo "$line" | cut -d: -f2 | sed 's/^ *//')
        elif [[ $line == "单核得分:"* ]]; then
            single=$(echo "$line" | cut -d: -f2 | sed 's/^ *//')
        elif [[ $line == "多核得分:"* ]]; then
            multi=$(echo "$line" | cut -d: -f2 | sed 's/^ *//')
        elif [[ $line == "浮点耗时:"* ]]; then
            float=$(echo "$line" | cut -d: -f2 | sed 's/^ *//' | awk '{print $1}')
            # 输出一行数据
            if [ -n "$vps_name" ] && [ -n "$single" ]; then
                printf "%-20s %-8s %-12s %-12s %-12s\n" "$vps_name" "$cores" "$single" "$multi" "${float}秒"
                # 重置变量
                vps_name=""; cores=""; single=""; multi=""; float=""
            fi
        fi
    done < "$COMPARE_FILE"
}

# 清理对比数据
clear_comparison() {
    if [ -f "$COMPARE_FILE" ]; then
        rm -f $COMPARE_FILE
        echo -e "${GREEN}已清空对比数据${NC}"
    else
        echo -e "${YELLOW}暂无对比数据可清理${NC}"
    fi
}

# 主菜单
main_menu() {
    while true; do
        echo ""
        echo -e "${BLUE}=== VPS性能测试工具 ===${NC}"
        echo "1. 运行性能测试"
        echo "2. 显示当前结果"
        echo "3. 添加到对比数据库"
        echo "4. 显示对比结果"
        echo "5. 生成对比表格"
        echo "6. 清空对比数据"
        echo "7. 显示性能参考"
        echo "8. 退出"
        echo ""
        echo -n "请选择操作 (1-8): "
        read -r choice
        
        case $choice in
            1)
                install_tools
                get_system_info
                run_benchmarks
                save_results
                show_current_results
                ;;
            2)
                if [ -f "$RESULTS_FILE" ]; then
                    show_current_results
                else
                    echo -e "${RED}请先运行性能测试${NC}"
                fi
                ;;
            3)
                if [ -f "$RESULTS_FILE" ]; then
                    add_comparison
                else
                    echo -e "${RED}请先运行性能测试${NC}"
                fi
                ;;
            4)
                show_comparison
                ;;
            5)
                generate_comparison_table
                ;;
            6)
                clear_comparison
                ;;
            7)
                show_reference
                ;;
            8)
                echo "再见!"
                exit 0
                ;;
            *)
                echo -e "${RED}无效选择${NC}"
                ;;
        esac
    done
}

# 直接运行测试并显示对比
if [ "$1" == "auto" ]; then
    install_tools
    get_system_info
    run_benchmarks
    show_current_results
    show_reference
    if [ -f "$COMPARE_FILE" ]; then
        generate_comparison_table
    fi
else
    main_menu
fi