#!/usr/bin/env bash

set -u
set -o pipefail

TEST_SIZE="${1:-20G}"
RUNTIME="${2:-300}"
MODE="${3:-safe}"

HOSTNAME="$(hostname)"
START_TIME="$(date +%Y%m%d_%H%M%S)"
LOG_DIR="./disk_auto_stress_${HOSTNAME}_${START_TIME}"

mkdir -p "$LOG_DIR"

apt install -y fio smartmontools lvm2

RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
NC="\033[0m"

log() {
    echo -e "$*" | tee -a "$LOG_DIR/run.log"
}

need_cmd() {
    if ! command -v "$1" >/dev/null 2>&1; then
        log "${RED}缺少命令: $1${NC}"
        return 1
    fi
    return 0
}

log "============================================================"
log "磁盘自动检测与压测脚本"
log "============================================================"
log "主机名: $HOSTNAME"
log "测试文件大小: $TEST_SIZE"
log "每项测试时长: ${RUNTIME}s"
log "模式: $MODE"
log "日志目录: $LOG_DIR"
log "============================================================"

if [[ "$EUID" -ne 0 ]]; then
    log "${YELLOW}建议使用 root 运行，否则 SMART 信息可能无法完整读取。${NC}"
fi

MISSING=0
for cmd in fio lsblk findmnt df awk sed grep sort uniq date hostname; do
    need_cmd "$cmd" || MISSING=1
done

if [[ "$MISSING" -eq 1 ]]; then
    log "${RED}依赖不完整，请先安装缺失命令。${NC}"
    log "Ubuntu/Debian 可执行:"
    log "  apt update && apt install -y fio smartmontools lvm2"
    exit 1
fi

if ! command -v smartctl >/dev/null 2>&1; then
    log "${YELLOW}警告: 未安装 smartctl，无法检查 SMART。${NC}"
    log "Ubuntu/Debian 可执行:"
    log "  apt update && apt install -y smartmontools"
fi

log ""
log "采集系统磁盘信息..."

lsblk -o NAME,TYPE,SIZE,FSTYPE,MOUNTPOINTS,MODEL,SERIAL > "$LOG_DIR/lsblk.log" 2>&1 || true
df -hT > "$LOG_DIR/df_hT.log" 2>&1 || true
findmnt -rno TARGET,SOURCE,FSTYPE,OPTIONS > "$LOG_DIR/findmnt.log" 2>&1 || true
mount > "$LOG_DIR/mount.log" 2>&1 || true

if command -v pvs >/dev/null 2>&1; then
    pvs > "$LOG_DIR/pvs.log" 2>&1 || true
fi

if command -v vgs >/dev/null 2>&1; then
    vgs > "$LOG_DIR/vgs.log" 2>&1 || true
fi

if command -v lvs >/dev/null 2>&1; then
    lvs -a -o +devices > "$LOG_DIR/lvs.log" 2>&1 || true
fi

log "${GREEN}基础磁盘信息已保存。${NC}"

get_parent_disk() {
    local dev="$1"
    local pkname=""

    if [[ -z "$dev" ]]; then
        return
    fi

    dev="$(readlink -f "$dev" 2>/dev/null || echo "$dev")"

    pkname="$(lsblk -no PKNAME "$dev" 2>/dev/null | head -n1 || true)"

    if [[ -n "$pkname" ]]; then
        echo "/dev/$pkname"
        return
    fi

    if [[ "$dev" =~ ^/dev/mapper/ ]] || [[ "$dev" =~ ^/dev/dm- ]]; then
        local slaves=""
        slaves="$(ls -1 "/sys/block/$(basename "$dev")/slaves" 2>/dev/null || true)"
        for s in $slaves; do
            local parent=""
            parent="$(lsblk -no PKNAME "/dev/$s" 2>/dev/null | head -n1 || true)"
            if [[ -n "$parent" ]]; then
                echo "/dev/$parent"
            else
                echo "/dev/$s"
            fi
        done
        return
    fi

    if [[ "$dev" =~ ^/dev/nvme[0-9]+n[0-9]+p[0-9]+$ ]]; then
        echo "$dev" | sed -E 's/p[0-9]+$//'
        return
    fi

    if [[ "$dev" =~ ^/dev/[a-z]+[0-9]+$ ]]; then
        echo "$dev" | sed -E 's/[0-9]+$//'
        return
    fi

    echo "$dev"
}

is_bad_mount() {
    local mp="$1"
    local src="$2"
    local fstype="$3"

    case "$fstype" in
        tmpfs|devtmpfs|overlay|squashfs|proc|sysfs|cgroup|cgroup2|securityfs|debugfs|tracefs|fusectl|configfs|efivarfs|autofs|mqueue|hugetlbfs|pstore|nsfs|ramfs)
            return 0
            ;;
    esac

    case "$mp" in
        /boot|/boot/efi|/run|/run/*|/dev|/dev/*|/proc|/proc/*|/sys|/sys/*)
            return 0
            ;;
    esac

    if [[ "$src" == overlay ]]; then
        return 0
    fi

    return 1
}

log ""
log "自动识别可测试挂载点..."

TEST_MOUNTS_FILE="$LOG_DIR/test_mounts.txt"
: > "$TEST_MOUNTS_FILE"

while read -r target source fstype; do
    [[ -z "${target:-}" ]] && continue

    if is_bad_mount "$target" "$source" "$fstype"; then
        continue
    fi

    case "$target" in
        /|/var/lib/docker|/data|/mnt/*|/opt|/opt/*)
            echo "$target $source $fstype" >> "$TEST_MOUNTS_FILE"
            ;;
    esac
done < <(findmnt -rno TARGET,SOURCE,FSTYPE)

if [[ ! -s "$TEST_MOUNTS_FILE" ]]; then
    log "${YELLOW}未识别到常规数据挂载点，默认测试根分区 /。${NC}"
    src="$(findmnt -rno SOURCE /)"
    fstype="$(findmnt -rno FSTYPE /)"
    echo "/ $src $fstype" >> "$TEST_MOUNTS_FILE"
fi

sort -u "$TEST_MOUNTS_FILE" -o "$TEST_MOUNTS_FILE"

log "将测试以下挂载点:"
cat "$TEST_MOUNTS_FILE" | tee -a "$LOG_DIR/run.log"

log ""
log "自动识别物理磁盘..."

PHYSICAL_DISKS_FILE="$LOG_DIR/physical_disks.txt"
: > "$PHYSICAL_DISKS_FILE"

while read -r target source fstype; do
    disk="$(get_parent_disk "$source" | head -n1 || true)"
    if [[ -n "$disk" && -b "$disk" ]]; then
        echo "$disk" >> "$PHYSICAL_DISKS_FILE"
    fi
done < "$TEST_MOUNTS_FILE"

# 兜底：列出所有 disk 类型设备
if [[ ! -s "$PHYSICAL_DISKS_FILE" ]]; then
    lsblk -dn -o NAME,TYPE | awk '$2=="disk"{print "/dev/"$1}' >> "$PHYSICAL_DISKS_FILE"
fi

sort -u "$PHYSICAL_DISKS_FILE" -o "$PHYSICAL_DISKS_FILE"

log "识别到的物理磁盘:"
cat "$PHYSICAL_DISKS_FILE" | tee -a "$LOG_DIR/run.log"

log ""
log "测试前 SMART 信息..."

if command -v smartctl >/dev/null 2>&1; then
    while read -r disk; do
        [[ -z "$disk" ]] && continue
        safe_name="$(basename "$disk")"
        log "采集 $disk SMART..."
        smartctl -a "$disk" > "$LOG_DIR/smart_before_${safe_name}.log" 2>&1 || true
        smartctl -H "$disk" > "$LOG_DIR/smart_health_before_${safe_name}.log" 2>&1 || true
    done < "$PHYSICAL_DISKS_FILE"
else
    log "${YELLOW}跳过 SMART 采集。${NC}"
fi

log ""
log "检查 dmesg 初始错误..."
dmesg -T > "$LOG_DIR/dmesg_before.log" 2>&1 || true

run_fio_test() {
    local mp="$1"
    local source="$2"
    local fstype="$3"

    local test_dir="${mp%/}/.disk_stress_test"
    if [[ "$mp" == "/" ]]; then
        test_dir="/.disk_stress_test"
    fi

    local tag
    tag="$(echo "$mp" | sed 's#/#_#g; s#^_##; s#_$##')"
    [[ -z "$tag" ]] && tag="root"

    log ""
    log "------------------------------------------------------------"
    log "开始测试挂载点: $mp"
    log "设备: $source"
    log "文件系统: $fstype"
    log "测试目录: $test_dir"
    log "------------------------------------------------------------"

    mkdir -p "$test_dir"

    df -hT "$mp" > "$LOG_DIR/df_before_${tag}.log" 2>&1 || true

    local avail_k
    avail_k="$(df -Pk "$mp" | awk 'NR==2{print $4}')"

    if [[ -z "$avail_k" || "$avail_k" -lt 1048576 ]]; then
        log "${YELLOW}可用空间不足 1GB，跳过 $mp。${NC}"
        return
    fi

    log "顺序写测试..."
    fio \
        --name="seq_write_${tag}" \
        --directory="$test_dir" \
        --filename="fio_test_file" \
        --size="$TEST_SIZE" \
        --rw=write \
        --bs=1M \
        --ioengine=libaio \
        --iodepth=32 \
        --direct=1 \
        --numjobs=1 \
        --runtime="$RUNTIME" \
        --time_based \
        --group_reporting \
        --eta=never \
        --output="$LOG_DIR/fio_${tag}_01_seq_write.log"

    log "顺序读测试..."
    fio \
        --name="seq_read_${tag}" \
        --directory="$test_dir" \
        --filename="fio_test_file" \
        --size="$TEST_SIZE" \
        --rw=read \
        --bs=1M \
        --ioengine=libaio \
        --iodepth=32 \
        --direct=1 \
        --numjobs=1 \
        --runtime="$RUNTIME" \
        --time_based \
        --group_reporting \
        --eta=never \
        --output="$LOG_DIR/fio_${tag}_02_seq_read.log"

    log "随机读写混合测试..."

    local iodepth="32"
    local numjobs="2"

    if [[ "$MODE" == "heavy" ]]; then
        iodepth="64"
        numjobs="4"
    fi

    fio \
        --name="randrw_${tag}" \
        --directory="$test_dir" \
        --filename="fio_test_file" \
        --size="$TEST_SIZE" \
        --rw=randrw \
        --rwmixread=70 \
        --bs=4k \
        --ioengine=libaio \
        --iodepth="$iodepth" \
        --direct=1 \
        --numjobs="$numjobs" \
        --runtime="$RUNTIME" \
        --time_based \
        --group_reporting \
        --eta=never \
        --output="$LOG_DIR/fio_${tag}_03_randrw.log"

    if [[ "$MODE" == "heavy" ]]; then
        log "heavy 模式：随机写测试..."
        fio \
            --name="randwrite_${tag}" \
            --directory="$test_dir" \
            --filename="fio_test_file" \
            --size="$TEST_SIZE" \
            --rw=randwrite \
            --bs=4k \
            --ioengine=libaio \
            --iodepth=64 \
            --direct=1 \
            --numjobs=4 \
            --runtime="$RUNTIME" \
            --time_based \
            --group_reporting \
            --eta=never \
            --output="$LOG_DIR/fio_${tag}_04_randwrite.log"
    fi

    df -hT "$mp" > "$LOG_DIR/df_after_${tag}.log" 2>&1 || true

    log "删除测试文件..."
    rm -f "$test_dir/fio_test_file"
    rmdir "$test_dir" 2>/dev/null || true

    log "${GREEN}$mp 测试完成。${NC}"
}

log ""
log "${YELLOW}即将开始 fio 压测。${NC}"
log "safe 模式: 顺序写、顺序读、随机混合读写"
log "heavy 模式: 额外增加随机写，压力更大"
log ""
log "如果这是生产机器，建议先用较小参数，例如:"
log "  $0 10G 120 safe"
log ""

sleep 3

while read -r target source fstype; do
    run_fio_test "$target" "$source" "$fstype"
done < "$TEST_MOUNTS_FILE"

log ""
log "测试后 SMART 信息..."

if command -v smartctl >/dev/null 2>&1; then
    while read -r disk; do
        [[ -z "$disk" ]] && continue
        safe_name="$(basename "$disk")"
        log "采集 $disk SMART..."
        smartctl -a "$disk" > "$LOG_DIR/smart_after_${safe_name}.log" 2>&1 || true
        smartctl -H "$disk" > "$LOG_DIR/smart_health_after_${safe_name}.log" 2>&1 || true
    done < "$PHYSICAL_DISKS_FILE"
fi

log ""
log "采集测试后 dmesg..."
dmesg -T > "$LOG_DIR/dmesg_after.log" 2>&1 || true

log ""
log "生成摘要..."

SUMMARY="$LOG_DIR/summary.txt"
: > "$SUMMARY"

{
    echo "============================================================"
    echo "磁盘压测摘要"
    echo "============================================================"
    echo "主机名: $HOSTNAME"
    echo "时间: $START_TIME"
    echo "测试文件大小: $TEST_SIZE"
    echo "每项测试时长: ${RUNTIME}s"
    echo "模式: $MODE"
    echo ""
    echo "测试挂载点:"
    cat "$TEST_MOUNTS_FILE"
    echo ""
    echo "物理磁盘:"
    cat "$PHYSICAL_DISKS_FILE"
    echo ""
    echo "============================================================"
    echo "fio 关键结果"
    echo "============================================================"
    grep -H -E "READ:|WRITE:|read: IOPS|write: IOPS|err=|lat \(|clat \(" "$LOG_DIR"/fio_*.log 2>/dev/null || true
    echo ""
    echo "============================================================"
    echo "SMART 健康状态"
    echo "============================================================"
    grep -H -E "SMART overall-health|SMART Health Status|PASSED|FAILED|Percentage Used|Media and Data Integrity Errors|Error Information Log Entries|Reallocated_Sector_Ct|Current_Pending_Sector|Offline_Uncorrectable|UDMA_CRC_Error_Count|Critical Warning|Available Spare|Data Units Written|Power On Hours|Unsafe Shutdowns" "$LOG_DIR"/smart_after_*.log 2>/dev/null || true
    echo ""
    echo "============================================================"
    echo "内核磁盘错误"
    echo "============================================================"
    grep -i -E "nvme|blk_update_request|I/O error|buffer I/O|reset|timeout|failed command|medium error|critical warning|ext4-fs error|xfs.*error|ata.*error|read-only file system" "$LOG_DIR/dmesg_after.log" 2>/dev/null || true
} > "$SUMMARY"

cat "$SUMMARY"

log ""
log "============================================================"
log "${GREEN}压测完成。${NC}"
log "日志目录: $LOG_DIR"
log "摘要文件: $SUMMARY"
log "============================================================"

log ""
log "如果发现以下情况，磁盘或链路就很可疑:"
log "1. fio 日志里出现 err 不为 0"
log "2. dmesg 出现 I/O error、timeout、reset、nvme critical warning"
log "3. SMART 里 Media and Data Integrity Errors 增加"
log "4. SMART 里 Error Information Log Entries 增加"
log "5. SATA 盘 UDMA_CRC_Error_Count 增加，通常还要检查线材或背板"