#!/usr/bin/env bash
# =============================================================================
# disk-health.sh — 硬盘健康度 / SMART / 读写量查询
# -----------------------------------------------------------------------------
# 功能:
#   * 列出所有物理磁盘(型号/容量/接口)
#   * SMART 健康状态 + 关键属性(温度/通电时长/重映射扇区等)
#   * 累计读取/写入量(SATA 走属性 241/242, NVMe 走 Data Units)
#   * 本次开机读写量(/sys/block stat)
# 用法:
#   bash disk-health.sh            # 默认: 概要查询
#   bash disk-health.sh --all      # 额外输出每块盘完整 smartctl -a
#   bash disk-health.sh --install-deps   # 安装 smartmontools(需要 root)
#   bash disk-health.sh --help
# 说明: SMART 需要 root 权限(smartctl); /sys 读写量普通用户可读。
#       虚拟盘(云主机常见)不支持 SMART 时自动降级为 /sys 统计。
# =============================================================================
set -u

MODE="summary"
ALL=0
HAVE_SMARTCTL=0
command -v smartctl >/dev/null 2>&1 && HAVE_SMARTCTL=1

for arg in "$@"; do
  case "$arg" in
    --all) ALL=1 ;;
    --install-deps) MODE="install" ;;
    --help|-h)
      sed -n '2,20p' "$0" | sed 's/^# \{0,1\}//'
      exit 0
      ;;
  esac
done

if [ "$MODE" = "install" ]; then
  if [ "$(id -u)" != "0" ]; then
    echo "需要 root 权限: sudo bash $0 --install-deps"
    exit 1
  fi
  if command -v apt-get >/dev/null 2>&1; then
    apt-get update -qq && apt-get install -y smartmontools
  elif command -v yum >/dev/null 2>&1; then
    yum install -y smartmontools
  elif command -v apk >/dev/null 2>&1; then
    apk add smartmontools
  else
    echo "不认识的包管理器，请手动安装 smartmontools"
    exit 1
  fi
  echo "安装完成，重新运行脚本即可"
  exit 0
fi

if [ "$(id -u)" != "0" ] && [ "$HAVE_SMARTCTL" = "1" ]; then
  echo "[!] 非 root 运行: SMART 属性可能读不到，/sys 读写量不受影响"
  echo
fi

human() {
  local bytes=$1
  awk -v b="$bytes" 'BEGIN{
    if (b < 0) { print "N/A"; exit }
    if (b < 1000)        printf "%d B\n", b
    else if (b < 1000000)      printf "%.1f KB\n", b/1000
    else if (b < 1000000000)   printf "%.1f MB\n", b/1000000
    else if (b < 1000000000000) printf "%.1f GB\n", b/1000000000
    else printf "%.2f TB\n", b/1000000000000
  }'
}

# 用 lsblk 拿容量(字节)
disk_size() {
  lsblk -dbn -o SIZE "/dev/$1" 2>/dev/null | head -1
}

# 本次开机读写(字节) — /sys/block/<dev>/stat 第3/7字段为扇区数
boot_rw() {
  local dev=$1 stat sectors_r sectors_w
  stat="/sys/block/$dev/stat"
  [ -r "$stat" ] || { echo "0 0"; return; }
  read -r sectors_r sectors_w < <(awk '{print $3, $7}' "$stat")
  echo "$((sectors_r * 512)) $((sectors_w * 512))"
}

smart_query() {
  local dev=$1
  smartctl -a "/dev/$dev" 2>/dev/null
}

# 输出一行健康度
report_disk() {
  local dev=$1 model="" size=""
  local health="" temp="" poh="" wear="" realloc="" pending="" uncorr=""
  local lread="" lwrite=""
  local brw boot_r boot_w
  local data

  if [ -r "/sys/block/$dev/device/model" ]; then
    model=$(tr -d ' ' < "/sys/block/$dev/device/model")
  fi
  size=$(disk_size "$dev")
  [ -n "$size" ] && size=$(human "$size")

  echo "===== /dev/$dev  ${model:-未知型号}  ${size:-?} ====="

  brw=$(boot_rw "$dev")
  boot_r=${brw%% *}
  boot_w=${brw##* }
  echo "本次开机读取: $(human "$boot_r")    本次开机写入: $(human "$boot_w")"

  if [ "$HAVE_SMARTCTL" = "1" ]; then
    data=$(smart_query "$dev")
    if [ -z "$data" ] || echo "$data" | grep -qi "Device does not support SMART\|Operation not supported\|Unknown device type"; then
      echo "SMART: 不可用(虚拟盘/不支持，属正常)"
    else
      # 通用健康状态
      health=$(echo "$data" | sed -n 's/^SMART overall-health self-assessment test result: *//p' | head -1)
      health=${health:-$(echo "$data" | grep -i "SMART Health Status" | awk -F': ' '{print $2}' | head -1)}
      [ -z "$health" ] && health="?"
      echo "SMART 状态: $health"

      # NVMe 专属
      temp=$(echo "$data" | sed -n 's/^Temperature: *//p' | head -1)
      wear=$(echo "$data" | sed -n 's/^Percentage Used: *//p' | head -1)
      poh=$(echo "$data" | sed -n 's/^Power On Hours: *//p' | head -1)
      crit=$(echo "$data" | sed -n 's/^Critical Warning: *//p' | head -1)
      [ -n "$crit" ] && echo "关键告警: $crit"

      # 累计读写 (NVMe)
      lread=$(echo "$data" | sed -n 's/^Data Units Read: *//p' | head -1)
      lwrite=$(echo "$data" | sed -n 's/^Data Units Written: *//p' | head -1)
      if [ -n "$lread" ]; then
        # 形如 "1,234,567 [632 GB]" 或 "1,234,567"
        lread_bytes=$(echo "$lread" | awk '{gsub(/,/,"",$1); print $1*512000}')
        lwrite_bytes=$(echo "$lwrite" | awk '{gsub(/,/,"",$1); print $1*512000}')
        echo "累计读取: $(human "$lread_bytes")    累计写入: $(human "$lwrite_bytes")"
      fi

      # SATA/ATA 属性表 (241=Total_LBAs_Written, 242=Total_LBAs_Read)
      if [ -z "$lread" ]; then
        lread_bytes=$(echo "$data" | awk '/^ *242 /{print $10}' | head -1)
        lwrite_bytes=$(echo "$data" | awk '/^ *241 /{print $10}' | head -1)
        if [ -n "${lread_bytes:-}" ]; then
          echo "累计读取: $(human "$((lread_bytes * 512))")    累计写入: $(human "$((lwrite_bytes * 512))")"
        fi
      fi

      # 温度 (ATA 属性 194)
      [ -z "$temp" ] && temp=$(echo "$data" | awk '/^ *194 /{print $10}' | head -1)
      # 通电时长 (属性 9)
      [ -z "$poh" ] && poh=$(echo "$data" | awk '/^ *9 /{print $10}' | head -1)
      # 坏块类属性
      realloc=$(echo "$data" | awk '/^ *(5|187) /{print $10}' | head -1)
      pending=$(echo "$data" | awk '/^ *197 /{print $10}' | head -1)
      uncorr=$(echo "$data" | awk '/^ *198 /{print $10}' | head -1)
      # NVMe 磨损用 Percentage Used(用量); SATA SSD 用 233/177 归一化值(剩余寿命)
      if [ -z "$wear" ]; then
        wear=$(echo "$data" | awk '/^ *233 /{print $4}' | head -1)
        [ -z "$wear" ] && wear=$(echo "$data" | awk '/^ *177 /{print $4}' | head -1)
        [ -n "$wear" ] && wear=$(echo "$wear" | sed 's/^0*//')
        [ -z "$wear" ] && wear=0
        wear="剩余 ${wear}%"
      else
        wear="已用 ${wear}"
      fi

      [ -n "$temp" ] && echo "温度: ${temp}°C"
      [ -n "$poh" ] && echo "通电时长: ${poh} 小时"
      [ -n "$wear" ] && echo "寿命: $wear"
      [ -n "$realloc" ] && echo "重映射扇区: $realloc"
      [ -n "$pending" ] && echo "待映射扇区: $pending"
      [ -n "$uncorr" ] && echo "不可纠正错误: $uncorr"
    fi

    if [ "$ALL" = "1" ]; then
      echo "--- 完整 smartctl 输出 ---"
      smartctl -a "/dev/$dev" 2>&1 | grep -vE "^===|^SMART Attributes|^ID#|^Vendor Specific|^Device Model|^Serial Number|^Firmware Version|^User Capacity|^Sector Sizes|^Rotation Rate|^Device is:|^ATA Version|^SATA Version|^Local Time|^SMART support|^General SMART|^SMART overall|^Copyright|^=== START"
      echo
    fi
  else
    echo "SMART: 未安装 smartctl(读写量仅本次开机)。root 跑 --install-deps 可安装"
  fi
  echo
}

main() {
  local disks
  echo "===== 硬盘健康报告  $(date '+%F %T')  ==="
  if command -v uptime >/dev/null 2>&1; then
    echo "系统已运行: $(uptime -p 2>/dev/null | sed 's/up //')（本次开机读写量为该时段的统计）"
    echo
  fi
  disks=$(lsblk -dn -o NAME,TYPE 2>/dev/null | awk '$2=="disk"{print $1}')
  if [ -z "$disks" ]; then
    echo "未找到物理磁盘"; exit 1
  fi
  for dev in $disks; do
    case "$dev" in loop*|ram*) continue ;; esac
    report_disk "$dev"
  done
}

main "$@"