#!/bin/bash
#====================================================================
# Description: Disk partition
# Author: Zhang JinWei
# Created: 2025-07-03
#====================================================================
set -e

# 日志函数
log() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1" >&2
}

# 检测磁盘设备
detect_disk_device() {
  for disk in vda sda; do
    [ -e "/dev/$disk" ] && echo "$disk" && return
  done
  log "错误: 未找到支持的磁盘设备"
  exit 1
}

# 检查是否使用LVM
check_lvm() {
  if ! command -v lvs &>/dev/null; then
    log "错误: LVM工具未安装"
    return 1
  fi
  lvs --noheadings &>/dev/null
  return $?
}

# 警告处理函数
warning_handler() {
  local command=$1
  set +e
  eval "${command}"
}

# 记录分区前后状态，找出新创建的分区
find_new_partition() {
  local disk=$1
  local before_file="/tmp/partitions_before.txt"
  local after_file="/tmp/partitions_after.txt"
  
  # 如果已经存在before文件，说明是在检查新分区
  if [ -f "${before_file}" ]; then
    lsblk -nlo NAME | grep "^${disk}[0-9]" > "${after_file}"
    # 找出新增的分区
    for part in $(comm -13 "${before_file}" "${after_file}"); do
      echo "${part}"
      rm -f "${before_file}" "${after_file}"
      return 0
    done
    rm -f "${before_file}" "${after_file}"
    return 1
  else
    # 记录当前分区状态
    lsblk -nlo NAME | grep "^${disk}[0-9]" > "${before_file}"
  fi
}

# 查找大于指定大小的未使用分区并扩容
find_unused_partition() {
  local disk=$1
  local min_size=$2  # 以GB为单位的最小大小
  local result=false
  
  log "查找大于${min_size}G的未使用分区..."
  
  # 获取所有分区
  for part in $(lsblk -nlo NAME | grep "^${disk}[0-9]"); do
    # 检查分区是否被挂载或添加到LVM
    if ! grep -q "/dev/${part}" /etc/mtab && ! pvs | grep -q "/dev/${part}"; then
      # 获取分区大小（以GB为单位）
      local part_size=$(lsblk -nlo SIZE -b /dev/${part} | awk '{printf "%.0f", $1/1024/1024/1024}')
      
      log "发现未使用分区: ${part}, 大小: ${part_size}G"
      
      # 检查分区大小是否大于指定值
      if [ "${part_size}" -ge "${min_size}" ]; then
        log "使用大于${min_size}G的未使用分区: ${part} 进行扩容"
        
        # 创建物理卷
        pvcreate "/dev/${part}" || {
          log "创建PV失败，尝试使用-ff选项"
          pvcreate -ff "/dev/${part}"
        }
        
        # 扩展卷组和逻辑卷
        vgextend ${VOLUME_GROUP} "/dev/${part}"
        lvextend -l +100%FREE /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
        
        result=true
        break
      fi
    fi
  done
  
  if [ "${result}" = true ]; then
    return 0
  else
    log "没有找到大于${min_size}G的未使用分区"
    return 1
  fi
}

# 检测所有额外磁盘
detect_additional_disks() {
  local system_disk=$1
  local disks=()

  for prefix in v s; do
    for letter in {a..z}; do
      local disk="${prefix}d${letter}"
      # 跳过系统盘
      [[ "${disk}" == "${system_disk}" ]] && continue
      
      # 检查磁盘是否存在且未被挂载
      if [[ -e "/dev/${disk}" ]] && ! grep -q "/dev/$disk" /etc/mtab; then
        # 检查磁盘是否已经被LVM使用
        if ! pvs | grep -q "/dev/${disk}"; then
          disks+=("${disk}")
        fi
      fi
    done
  done

  echo "${disks[@]:-}"
}

# 修复分区表
fix_partition_table() {
  fdisk -l 2>/tmp/status
  if [ -s "/tmp/status" ]; then
    /usr/bin/expect <<EOF
spawn parted -l
expect {
    "OK/Cancel?" { send "OK\r"; exp_continue }
    "Fix/Ignore?" { send "Fix\r"; exp_continue }
    timeout { exit 1 }
    eof
}
EOF
  fi
  rm -f /tmp/status
}

# 扩展文件系统
extend_filesystem() {
  local fs_type=$1
  log "开始扩展文件系统..."
  if [[ "${fs_type}" == "xfs" ]]; then
    xfs_growfs /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
  else
    resize2fs /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
  fi
}

# 处理主磁盘
process_system_disk() {
  local disk=$1
  
  # 记录当前分区状态
  find_new_partition "${disk}"
  
  log "主磁盘创建新分区进行扩容"
  # 创建新分区
  warning_handler "echo -e 'n\np\n\n\n\nw' | fdisk /dev/${disk}"
  
  # 更新分区表
  partprobe /dev/${disk}
  sleep 2
  
  # 查找新创建的分区
  local NEW_PART=""
  local retries=0
  while [ -z "${NEW_PART}" ] && [ ${retries} -lt 10 ]; do
    NEW_PART=$(find_new_partition "${disk}")
    if [ -z "${NEW_PART}" ]; then
      log "等待新分区出现..."
      sleep 1
      retries=$((retries + 1))
    fi
  done
  
  if [ -n "${NEW_PART}" ]; then
    log "找到新分区: ${NEW_PART}"
    pvcreate "/dev/${NEW_PART}" || {
      log "创建PV失败，尝试使用-ff选项"
      pvcreate -ff "/dev/${NEW_PART}"
    }
    vgextend ${VOLUME_GROUP} "/dev/${NEW_PART}"
    lvextend -l +100%FREE /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
    return 0
  else
    log "未能找到新创建的分区，扩容失败"
    return 1
  fi
}

# 处理额外磁盘
process_additional_disk() {
  local disk=$1
  log "处理额外磁盘: ${disk}"

  # 检查磁盘是否已经包含分区
  if lsblk | grep -q "${disk}[0-9]"; then
    # 尝试使用现有未使用分区
    for partition in $(lsblk -nlo NAME | grep "^${disk}[0-9]"); do
      if ! grep -q "/dev/${partition}" /etc/mtab && ! pvs | grep -q "/dev/${partition}"; then
        log "使用现有分区 ${partition} 进行LVM扩展"
        pvcreate "/dev/${partition}"
        vgextend "${VOLUME_GROUP}" "/dev/${partition}"
        lvextend -l +100%FREE "/dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}"
        return 0
      fi
    done
    log "磁盘 ${disk} 上没有可用分区"
    return 1
  fi

  # 创建新分区
  log "在磁盘 ${disk} 上创建新分区"
  warning_handler "echo -e 'n\np\n\n\n\nw' | fdisk /dev/${disk}"

  partprobe "/dev/${disk}"
  sleep 2

  # 将新分区添加到LVM
  log "将 ${disk}1 添加到LVM"
  pvcreate "/dev/${disk}1"
  vgextend "${VOLUME_GROUP}" "/dev/${disk}1"
  lvextend -l +100%FREE "/dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}"

  return 0
}

# 主函数
main() {
  # 获取系统信息
  SYSTEM_DISK=$(detect_disk_device)
  TOTAL_DISK_SIZE=$(lsblk | awk '/^'"${SYSTEM_DISK}"'/{print $4}')
  
  # 检查是否使用LVM
  if ! check_lvm; then
    log "错误: 系统未使用LVM，无法扩容"
    exit 1
  fi
  
  # 获取LVM信息
  VOLUME_GROUP=$(vgs --noheadings -o vg_name | head -1 | tr -d '[:space:]')
  LOGICAL_VOLUME=$(lvs --noheadings -o lv_name | head -1 | tr -d '[:space:]')
  
  if [ -z "${VOLUME_GROUP}" ] || [ -z "${LOGICAL_VOLUME}" ]; then
    log "错误: 未能获取LVM信息"
    exit 1
  fi
  
  # 检查系统盘大小
  if [[ "${TOTAL_DISK_SIZE%G}" -le "20" ]]; then
    log "磁盘大小小于或等于20G，不需要扩容"
    exit 0
  else
    log "系统磁盘大于20G，开始进行分区扩展"
    
    # 修复分区表
    fix_partition_table
    
    # 先尝试查找并使用大于20G的未使用分区
    if ! find_unused_partition "${SYSTEM_DISK}" 20; then
      # 如果没有找到合适的未使用分区，则创建新分区
      log "没有找到合适的未使用分区，将创建新分区"
      
      # 处理系统盘
      process_system_disk "${SYSTEM_DISK}"
    fi
    
    # 获取文件系统类型
    FS_TYPE=$(grep -Ev '^$|^#' /etc/fstab | awk '/\// && !/boot/{print $3}' | head -1)
    
    # 扩展文件系统
    extend_filesystem "${FS_TYPE}"
  fi
  
  # 检测额外磁盘
  ADDITIONAL_DISKS=($(detect_additional_disks "${SYSTEM_DISK}"))
  
  # 处理额外磁盘
  if [ ${#ADDITIONAL_DISKS[@]} -gt 0 ]; then
    log "发现额外磁盘: ${ADDITIONAL_DISKS[*]}"
    
    for disk in "${ADDITIONAL_DISKS[@]}"; do
      process_additional_disk "${disk}"
    done
    
    # 再次扩展文件系统
    extend_filesystem "${FS_TYPE}"
  fi
  
  log "磁盘扩容完成"
}

# 执行主函数
main "$@"
