#!/usr/bin/env bash
set -euo pipefail

###############################################################################
# FileWave Smart Disk Expansion (filewave-expand-disk.sh)
# Features:
# - LVM byte outputs forced to integers (no scientific notation)
# - PV device list deduplicated
# - Clean terminal output by default; --verbose shows command output
# - Smart tool detection and installation
###############################################################################

DRY_RUN=0
REBOOT=0
VERBOSE=0

usage() {
  cat <<'EOF'
Usage: sudo ./filewave-expand-disk.sh [options]

Options:
  --dry-run      Print planned actions; do not modify anything
  --reboot       Reboot at the end
  --verbose      Show command output in terminal (still logs everything)
  -h, --help     Show help

Notes:
- Must be run as root.
- Log file is written to /var/log and made world-readable.
- Required tools will be installed automatically if missing.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --dry-run)  DRY_RUN=1; shift ;;
    --reboot)   REBOOT=1; shift ;;
    --verbose)  VERBOSE=1; shift ;;
    -h|--help)  usage; exit 0 ;;
    *)
      echo "ERROR: Unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

if [[ "${EUID}" -ne 0 ]]; then
  cat >&2 <<'EOF'
ERROR: This script must be run as root.

Please re-run using sudo:
  sudo ./filewave-expand-disk.sh [options]
EOF
  exit 1
fi

ts="$(date -u +%Y%m%d-%H%M%S)"
LOG_FILE="/var/log/filewave-expand-disk-${ts}.log"
touch "$LOG_FILE"
chmod 0644 "$LOG_FILE"

echo "Log file: $LOG_FILE"
echo "If opening a support case, attach: $LOG_FILE"
echo

log() {
  printf '[%s] %s\n' "$(date -u +%FT%TZ)" "$*" | tee -a "$LOG_FILE" >/dev/null
}

run() {
  log "+ $*"
  if [[ "$DRY_RUN" -eq 1 ]]; then
    log "DRY_RUN: command not executed"
    return 0
  fi

  if [[ "$VERBOSE" -eq 1 ]]; then
    "$@" 2>&1 | tee -a "$LOG_FILE"
  else
    "$@" >>"$LOG_FILE" 2>&1
  fi
}

section() {
  echo "== $1 =="
  echo
  log "== $1 =="
}

human_bytes() {
  local b="${1:-0}"
  if command -v numfmt >/dev/null 2>&1; then
    numfmt --to=iec --suffix=B "$b" 2>/dev/null || echo "${b}B"
  else
    echo "${b}B"
  fi
}

# Force LVM to give integer bytes (no suffix, no sci-notation)
lvm_bytes() {
  # Usage: lvm_bytes "vgs args..." OR "lvs args..."
  # Must return an integer or 0.
  local out
  out="$("$@" 2>/dev/null | tr -d ' ' | head -n1 || true)"
  # empty -> 0
  [[ -z "$out" ]] && { echo 0; return; }
  # defensive: strip anything non-digit (should already be digits)
  echo "$out" | tr -cd '0-9'
}

section "Checking and installing prerequisites"
log "dry_run=$DRY_RUN reboot=$REBOOT verbose=$VERBOSE log_file=$LOG_FILE"

# Check for required tools and install if missing
check_and_install_tools() {
  local missing=()
  
  # Check each required command
  for cmd in growpart pvresize lvextend resize2fs xfs_growfs sgdisk partx blockdev udevadm lsblk lvs vgs df findmnt numfmt; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
      missing+=("$cmd")
    fi
  done
  
  if [[ ${#missing[@]} -gt 0 ]]; then
    echo "Missing required tools: ${missing[*]}"
    log "Missing tools: ${missing[*]}. Installing prerequisite packages..."
    
    run apt-get update -y
    # cloud-guest-utils provides growpart
    # lvm2 provides lvs, vgs, lvextend, pvresize
    # util-linux provides blockdev, lsblk, findmnt
    # gdisk provides sgdisk
    # e2fsprogs provides resize2fs
    # xfsprogs provides xfs_growfs
    run apt-get install -y cloud-guest-utils lvm2 util-linux gdisk e2fsprogs xfsprogs
    
    log "Tool installation complete. Verifying..."
    local still_missing=()
    for cmd in growpart pvresize lvextend; do
      if ! command -v "$cmd" >/dev/null 2>&1; then
        still_missing+=("$cmd")
      fi
    done
    
    if [[ ${#still_missing[@]} -gt 0 ]]; then
      echo "ERROR: Critical tools still missing after installation: ${still_missing[*]}" >&2
      log "ERROR: Critical tools still missing: ${still_missing[*]}"
      exit 1
    fi
  else
    log "All required tools are present; skipping installation."
  fi
}

check_and_install_tools
echo

section "Detecting root LVM layout"

ROOT_SRC="$(findmnt -n -o SOURCE / || true)"
if [[ -z "$ROOT_SRC" ]]; then
  echo "ERROR: Unable to determine root filesystem source." >&2
  log "ERROR: Unable to determine root SOURCE via findmnt."
  exit 1
fi

# Resolve root LV
ROOT_LV_PATH=""
VG_NAME=""
LV_NAME=""

if lvs --noheadings "$ROOT_SRC" >/dev/null 2>&1; then
  read -r VG_NAME LV_NAME < <(lvs --noheadings -o vg_name,lv_name "$ROOT_SRC" | awk '{print $1, $2}')
  ROOT_LV_PATH="/dev/${VG_NAME}/${LV_NAME}"
else
  echo "ERROR: Root filesystem ($ROOT_SRC) is not an LVM logical volume; this script targets LVM appliances." >&2
  log "ERROR: Root filesystem ($ROOT_SRC) not recognized as LVM LV."
  exit 1
fi

FSTYPE="$(findmnt -n -o FSTYPE / || true)"

echo "Root LV: $ROOT_LV_PATH"
echo "VG:      $VG_NAME"
echo "FS type: ${FSTYPE:-unknown}"
echo
log "root_source=$ROOT_SRC root_lvpath=$ROOT_LV_PATH vg_name=$VG_NAME lv_name=$LV_NAME root_fstype=$FSTYPE"

section "Finding PV backing the root LV"

# Pull devices backing LV; strip "(...)" and commas; dedupe
LV_DEVS="$(
  lvs --noheadings -o devices "$ROOT_LV_PATH" \
  | awk '{print $1}' \
  | sed 's/([^)]*)//g' \
  | tr -d ',' \
  | awk 'NF' \
  | sort -u \
  | tr '\n' ' '
)"
LV_DEVS="$(echo "$LV_DEVS" | xargs || true)"

if [[ -z "$LV_DEVS" ]]; then
  echo "ERROR: Unable to determine PV device(s) behind $ROOT_LV_PATH." >&2
  log "ERROR: Unable to determine LV devices."
  exit 1
fi

echo "PV(s):   $LV_DEVS"
echo
log "lv_devices=$LV_DEVS"

vg_free_bytes() {
  lvm_bytes vgs --noheadings --nosuffix --units b -o vg_free "$VG_NAME"
}
lv_size_bytes() {
  lvm_bytes lvs --noheadings --nosuffix --units b -o lv_size "$ROOT_LV_PATH"
}
fs_size_bytes() { df --output=size -B1 / | tail -n1 | tr -d ' '; }
fs_avail_bytes(){ df --output=avail -B1 / | tail -n1 | tr -d ' '; }

section "BEFORE state"
BEFORE_VG_FREE="$(vg_free_bytes)"
BEFORE_LV="$(lv_size_bytes)"
BEFORE_FS="$(fs_size_bytes)"
BEFORE_AVAIL="$(fs_avail_bytes)"

echo "VG free:      $(human_bytes "$BEFORE_VG_FREE")"
echo "LV size:      $(human_bytes "$BEFORE_LV")"
echo "FS size (/):  $(human_bytes "$BEFORE_FS")"
echo "FS avail (/): $(human_bytes "$BEFORE_AVAIL")"
echo

log "before_vg_free_bytes=$BEFORE_VG_FREE"
log "before_lv_bytes=$BEFORE_LV"
log "before_fs_bytes=$BEFORE_FS"
log "before_fs_avail_bytes=$BEFORE_AVAIL"

detect_parent_sysfs() {
  local dev="$1"
  local real
  real="$(readlink -f "$dev" 2>/dev/null || echo "$dev")"

  DET_KIND="unknown"
  DET_PARENT="$real"
  DET_PARTNUM=""

  local type pkname
  type="$(lsblk -n -o TYPE "$real" 2>/dev/null | head -n1 || true)"
  pkname="$(lsblk -n -o PKNAME "$real" 2>/dev/null | head -n1 || true)"

  if [[ "$type" == "part" && -n "$pkname" ]]; then
    DET_KIND="part"
    DET_PARENT="/dev/$pkname"
    DET_PARTNUM="$(echo "$real" | sed -E 's@.*/[^0-9]*([0-9]+)$@\1@' || true)"
  elif [[ "$type" == "disk" ]]; then
    DET_KIND="disk"
    DET_PARENT="$real"
  fi

  log "detect_parent_sysfs: dev=$dev real=$real kind=$DET_KIND parent=$DET_PARENT partnum=${DET_PARTNUM:-none}"
}

STRONGLY_RECOMMEND_REBOOT=0
reboot_recommend_reason=""

try_gpt_fix_and_rescan() {
  local disk="$1"
  log "Rescanning disk (non-interactive, best-effort): $disk"

  local pttype
  pttype="$(lsblk -n -o PTTYPE "$disk" 2>/dev/null | head -n1 || true)"
  if [[ "$pttype" == "gpt" ]]; then
    local out
    out="$({ sgdisk -e "$disk"; } 2>&1)" || true
    log "+ sgdisk -e $disk"
    printf "%s\n" "$out" >>"$LOG_FILE"

    if echo "$out" | grep -qi "kernel is still using the old partition table"; then
      STRONGLY_RECOMMEND_REBOOT=1
      reboot_recommend_reason="Kernel is still using the old partition table after GPT repair (sgdisk -e)."
      log "WARN: kernel still using old partition table; partx/udevadm may refresh; otherwise reboot."
    fi
  else
    log "PTTYPE for $disk is '$pttype' (not gpt); skipping sgdisk -e."
  fi

  run partx -u "$disk" || true

  local bd_out bd_rc
  bd_out="$({ blockdev --rereadpt "$disk"; } 2>&1)" && bd_rc=0 || bd_rc=$?
  if [[ "$bd_rc" -ne 0 ]]; then
    if echo "$bd_out" | grep -qi "Device or resource busy"; then
      log "WARN: blockdev --rereadpt reported busy (expected on mounted/in-use disks). Continuing."
      log "blockdev_output=$bd_out"
    else
      log "WARN: blockdev --rereadpt failed: $bd_out"
    fi
  else
    log "+ blockdev --rereadpt $disk"
    printf "%s\n" "$bd_out" >>"$LOG_FILE"
  fi

  run udevadm settle || true
}

section "Expanding PV(s) and partition(s) (if needed)"

declare -A GROW_STATUS
ANY_GROWPART_CHANGED=0

for pv in $LV_DEVS; do
  echo "Processing: $pv"
  log "Processing PV backing device: $pv"

  detect_parent_sysfs "$pv"
  local_parent="$DET_PARENT"
  local_partnum="${DET_PARTNUM:-}"
  local_kind="$DET_KIND"

  if [[ "$local_kind" == "part" && -n "$local_partnum" ]]; then
    try_gpt_fix_and_rescan "$local_parent"

    if [[ "$DRY_RUN" -eq 1 ]]; then
      echo "  - would run: growpart $local_parent $local_partnum"
      log "DRY_RUN: would run growpart $local_parent $local_partnum"
      GROW_STATUS["$pv"]="SKIPPED"
    else
      gp_out="$({ growpart "$local_parent" "$local_partnum"; } 2>&1)" || true
      printf "%s\n" "$gp_out" >>"$LOG_FILE"

      if echo "$gp_out" | grep -q "^CHANGED:"; then
        echo "  - action: grew partition ${local_parent} ${local_partnum} (growpart)"
        log "growpart_status__${pv}=CHANGED"
        GROW_STATUS["$pv"]="CHANGED"
        ANY_GROWPART_CHANGED=1
      elif echo "$gp_out" | grep -q "^NOCHANGE:"; then
        log "growpart_status__${pv}=NOCHANGE"
        GROW_STATUS["$pv"]="NOCHANGE"
      else
        log "growpart_status__${pv}=UNKNOWN"
        GROW_STATUS["$pv"]="UNKNOWN"
      fi
    fi

    echo "  - pvresize $pv"
    run pvresize "$pv"

  elif [[ "$local_kind" == "disk" ]]; then
    GROW_STATUS["$pv"]="WHOLEDISK"
    echo "  - pvresize $pv"
    run pvresize "$pv"
  else
    GROW_STATUS["$pv"]="UNKNOWN"
    echo "  - pvresize $pv"
    run pvresize "$pv"
  fi
done

VG_FREE_AFTER_PV="$(vg_free_bytes)"
log "vg_free_bytes_after_pvresize=$VG_FREE_AFTER_PV"
echo
echo "VG free after pvresize: $(human_bytes "$VG_FREE_AFTER_PV")"
echo

if [[ "$VG_FREE_AFTER_PV" -gt 0 ]]; then
  echo "Action: detected free space in VG after pvresize ($(human_bytes "$VG_FREE_AFTER_PV"))."
  echo
fi

if [[ "$STRONGLY_RECOMMEND_REBOOT" -eq 1 ]]; then
  echo "RECOMMENDATION: Reboot strongly recommended."
  echo "Reason: $reboot_recommend_reason"
  echo "Tip: You can re-run this script after reboot to ensure all space is consumed."
  echo
  log "RECOMMEND_REBOOT_STRONG=1 reason=$reboot_recommend_reason"
fi

ALL_NONEXPANDING=1
for pv in $LV_DEVS; do
  st="${GROW_STATUS[$pv]:-UNKNOWN}"
  [[ "$st" == "CHANGED" ]] && ALL_NONEXPANDING=0
done

NOOP_SHORT_CIRCUIT=0
if [[ "$VG_FREE_AFTER_PV" -le 0 && "$ALL_NONEXPANDING" -eq 1 ]]; then
  NOOP_SHORT_CIRCUIT=1
fi

if [[ "$NOOP_SHORT_CIRCUIT" -eq 1 ]]; then
  section "No expansion needed"
  echo "Disk/partition appears fully expanded and VG has no free space."
  echo
  log "No-op short-circuit: no growpart CHANGED and VG has 0 free bytes; skipping lvextend + filesystem grow."
fi

section "Extending root LV and filesystem"

if [[ "$NOOP_SHORT_CIRCUIT" -eq 1 ]]; then
  log "Skipping LV/filesystem extension due to no-op short-circuit."
else
  echo "Action: extending root LV to consume free VG space..."
  run lvextend -l +100%FREE "$ROOT_LV_PATH"

  case "${FSTYPE}" in
    ext4|ext3|ext2)
      echo "Action: growing ext filesystem online (resize2fs)..."
      run resize2fs "$ROOT_LV_PATH"
      ;;
    xfs)
      echo "Action: growing XFS filesystem online (xfs_growfs /)..."
      run xfs_growfs /
      ;;
    *)
      echo "WARN: Unsupported/unknown filesystem type '$FSTYPE'; skipping filesystem grow."
      log "WARN: Unsupported/unknown filesystem type '$FSTYPE'; skipping filesystem grow."
      ;;
  esac
  echo
fi

section "AFTER state"

AFTER_VG_FREE="$(vg_free_bytes)"
AFTER_LV="$(lv_size_bytes)"
AFTER_FS="$(fs_size_bytes)"
AFTER_AVAIL="$(fs_avail_bytes)"

echo "VG free:      $(human_bytes "$AFTER_VG_FREE")"
echo "LV size:      $(human_bytes "$AFTER_LV")"
echo "FS size (/):  $(human_bytes "$AFTER_FS")"
echo "FS avail (/): $(human_bytes "$AFTER_AVAIL")"
echo

log "after_vg_free_bytes=$AFTER_VG_FREE"
log "after_lv_bytes=$AFTER_LV"
log "after_fs_bytes=$AFTER_FS"
log "after_fs_avail_bytes=$AFTER_AVAIL"

DELTA_LV=$(( AFTER_LV - BEFORE_LV ))
DELTA_FS=$(( AFTER_FS - BEFORE_FS ))

if [[ "$DELTA_LV" -gt 0 || "$DELTA_FS" -gt 0 ]]; then
  msg="SUMMARY: Expansion performed. LV grew by $(human_bytes "$DELTA_LV") and filesystem grew by $(human_bytes "$DELTA_FS")."
  echo "$msg"
  log "$msg"
else
  msg="SUMMARY: No measurable growth detected. This is normal if the disk was already fully expanded."
  echo "$msg"
  log "WARN: $msg"
fi
echo

section "Post-operation state (for support)"
run lsblk
run pvs
run vgs
run lvs
run df -h /

echo
echo "Support log file: $LOG_FILE"
echo
log "Support log file: $LOG_FILE"

if [[ "$REBOOT" -eq 1 ]]; then
  echo "Reboot requested (--reboot). Rebooting now..."
  log "Reboot requested. Rebooting now..."
  run reboot
else
  echo "Reboot not requested (default). No reboot will be performed."
  if [[ "$STRONGLY_RECOMMEND_REBOOT" -eq 1 ]]; then
    echo "NOTE: A reboot is strongly recommended due to partition table refresh warnings."
  fi
  log "Reboot not requested (default). No reboot will be performed."
fi