#!/bin/bash
#====================================================================
# Description: Disk partition
# Author: Liu YanHao
# Created: 2024-11-9
# Modified: 2024-12-12
#====================================================================
set -euo pipefail

# Define logging function
log() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1"
}

# Get system information, use readonly to prevent accidental modification
detect_disk_device() {
  for disk in vda sda; do
    [ -e "/dev/$disk" ] && echo "$disk" && return
  done
  log "Error: No supported disk device found"
  exit 1
}

readonly SYSTEM_DISK=$(detect_disk_device)
readonly PARTITION_COUNT=$(lsblk | grep -c "${SYSTEM_DISK}")
readonly VOLUME_GROUP=$(vgs --noheadings -o vg_name | head -1 | tr -d '[:space:]')
readonly LOGICAL_VOLUME=$(lvs --noheadings -o lv_name | head -1 | tr -d '[:space:]')
readonly TOTAL_DISK_SIZE=$(lsblk | awk '/^'"${SYSTEM_DISK}"'/{print $4}')

# Check if required commands exist
for cmd in fdisk parted pvcreate vgextend lvextend; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    log "Error: Command not found - $cmd"
    exit 1
  fi
done

# Optimize warning handler function
warning_handler() {
  local command=$1
  set +e
  eval "${command}"
}

# Fix partition table
fix_partition_table() {
  # Check and install expect
  if ! command -v expect >/dev/null 2>&1; then
    (yum install -y expect || dnf install -y expect) || (apt-get update && apt-get install -y expect) || {
      log "Error: Unsupported OS for expect installation"
      exit 1
    }
  fi

  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
extend_filesystem() {
  local fs_type=$1
  log "Starting filesystem extension..."
  if [[ "${fs_type}" == "xfs" ]]; then
    xfs_growfs /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
  else
    resize2fs /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}
  fi
}

# Main logic
main() {
  # Load OS information
  if [ -f /etc/os-release ]; then
    . /etc/os-release
  else
    log "Error: Cannot find /etc/os-release file"
    exit 1
  fi

  # Ensure VERSION_ID is set
  if [ -z "${VERSION_ID+x}" ]; then
    log "Error: VERSION_ID is not set"
    exit 1
  fi

  # Check disk size
  if [[ "${TOTAL_DISK_SIZE%G}" -le "20" ]]; then
    log "Disk size is less than 20G, no extension needed"
    exit 0
  fi

  # Fix partition table
  fix_partition_table

  # Execute different partition commands based on the operating system
  case $(
    . /etc/os-release
    echo "${ID}${VERSION_ID}"
  ) in
  "centos9")
    log "Processing CentOS Stream 9"
    CentOS_Stream_9
    ;;
  "centos8")
    log "Processing CentOS Stream 8"
    CentOS_Stream_8
    ;;
  "centos7"*)
    log "Processing CentOS 7"
    warning_handler "echo -e 'n\np\n\n\n\nw' | fdisk /dev/${SYSTEM_DISK}"
    ;;
  "ubuntu24.04")
    log "Processing Ubuntu 24.04"
    ;;
  "ubuntu"*)
    log "Processing Ubuntu"
    if [[ "${VERSION_ID}" == "16.04" ]]; then
      warning_handler "echo -e 'n\np\n\n\n\nw' | fdisk /dev/${SYSTEM_DISK}"
    else
      warning_handler "echo -e 'n\n\n\n\nw' | fdisk /dev/${SYSTEM_DISK}"
    fi
    ;;
  *)
    log "Unsupported operating system"
    exit 1
    ;;
  esac

  # Wait for new partition to appear
  local timeout=30
  local count=0
  while ! lsblk | grep -qi "${SYSTEM_DISK}${PARTITION_COUNT}"; do
    if ((count >= timeout)); then
      log "Timeout waiting for new partition"
      exit 1
    fi
    partprobe /dev/${SYSTEM_DISK}
    sleep 1
    ((count++))
  done

  # Extend logical volume
  log "Creating physical volume and extending logical volume..."
  pvcreate /dev/${SYSTEM_DISK}${PARTITION_COUNT}
  vgextend ${VOLUME_GROUP} /dev/${SYSTEM_DISK}${PARTITION_COUNT}
  lvextend -l +100%FREE /dev/${VOLUME_GROUP}/${LOGICAL_VOLUME}

  # Extend filesystem
  readonly File_System=$(grep -Ev '^$|^#' /etc/fstab | awk '/boot/{print $3}')
  extend_filesystem "${File_System}"

  log "Disk extension completed"
}

# Execute main function
main "$@"

