#!/bin/bash

# 1. 自动安装依赖
apt update && apt install -y fio jq

echo "============================================"
echo "          磁盘性能自定义测试工具            "
echo "============================================"

# 2. 交互输入路径
read -p "请输入测试路径 (设备如 /dev/vdb 或 目录如 /root): " USER_PATH

# 3. 交互输入测试大小
read -p "请输入测试文件大小 (例如 1G, 512M, 10G): " TEST_SIZE
[[ -z "$TEST_SIZE" ]] && TEST_SIZE="1G" # 默认 1G

# 4. 判断路径类型并设置参数
if [ -b "$USER_PATH" ]; then
    FILENAME_OPTS="--filename=$USER_PATH"
    echo "警告：原始设备测试将擦除 $USER_PATH 的所有数据！"
    read -p "确认继续？(y/n): " CONFIRM
    [[ "$CONFIRM" != "y" ]] && exit 1
elif [ -d "$USER_PATH" ]; then
    # 交互输入自定义文件名
    read -p "请输入临时测试文件名 (默认 fio_test.bin): " CUSTOM_FILE
    [[ -z "$CUSTOM_FILE" ]] && CUSTOM_FILE="fio_test.bin"
    
    FILENAME_OPTS="--filename=$USER_PATH/$CUSTOM_FILE"
    TEST_FILE_PATH="$USER_PATH/$CUSTOM_FILE"
    echo "模式：安全模式（在 $USER_PATH 下创建 $TEST_SIZE 测试文件）"
else
    echo "错误：路径无效！" && exit 1
fi

# 5. 运行测试
echo "正在进行 $TEST_SIZE 的读写测试，请稍后..."
fio --ioengine=libaio --direct=1 --group_reporting \
    $FILENAME_OPTS --size=$TEST_SIZE --runtime=60 \
    --name=Read_Test --rw=read --bs=1M \
    --name=Write_Test --rw=write --bs=1M \
    --output-format=json > res.json

# 6. 数据解析 (增加单位转换处理)
R_BW_KB=$(jq '.jobs[] | select(.jobname=="Read_Test") | .read.bw_mean' res.json 2>/dev/null || echo 0)
W_BW_KB=$(jq '.jobs[] | select(.jobname=="Write_Test") | .write.bw_mean' res.json 2>/dev/null || echo 0)

# 7. 显示结果
echo "------------------------------------"
if [ "$R_BW_KB" -eq 0 ] && [ "$W_BW_KB" -eq 0 ]; then
    echo "测试失败：请检查权限或路径是否正确。"
else
    # 转换为 MB/s 并保留一位小数（使用 awk 处理浮点）
    R_SPEED=$(awk "BEGIN {printf \"%.2f\", $R_BW_KB/1024}")
    W_SPEED=$(awk "BEGIN {printf \"%.2f\", $W_BW_KB/1024}")
    echo "测试路径: $USER_PATH"
    echo "样本大小: $TEST_SIZE"
    echo "平均读取速度: $R_SPEED MB/s"
    echo "平均写入速度: $W_SPEED MB/s"
fi
echo "------------------------------------"

# 8. 清理临时文件 (仅限目录模式)
if [ -f "$TEST_FILE_PATH" ]; then
    read -p "测试完成，是否删除临时文件 $CUSTOM_FILE? (y/n): " DEL_CONFIRM
    [[ "$DEL_CONFIRM" == "y" ]] && rm "$TEST_FILE_PATH" && echo "临时文件已清理。"
fi