#!/bin/bash
#Nifty: Safely extend /var/www and / on NixFactory web servers
# Extends:
#   - /var/www (ext4 on /dev/mapper/datavg-lv0 via PV /dev/sdc)
#   - /       (xfs on /dev/mapper/vg-root via PV /dev/sda2) OR plain / on /dev/sda4 (xfs/ext4)
#
# Policy:
#   - Only proceed on LVM targets if we would gain at least MIN_ADD_PE full extents.
#     Example: with 4 MiB PE size, MIN_ADD_PE=256 ≈ 1 GiB minimum effective growth.

set -o pipefail
set -u
export LVM_SUPPRESS_FD_WARNINGS=1

# ---------- Colors & logging ----------
RED="\e[31m"; GRN="\e[32m"; YLW="\e[33m"; BLU="\e[34m"; NRM="\e[39m"
log()  { printf "%b[INFO]%b %s\n"  "$BLU" "$NRM" "$*"; }
ok()   { printf "%b[OK]%b   %s\n"  "$GRN" "$NRM" "$*"; }
warn() { printf "%b[WARN]%b %s\n" "$YLW" "$NRM" "$*"; }
err2() { printf "%b[FAIL]%b %s\n" "$RED" "$NRM" "$*"; }  # avoid clobbering niftylib's err()

# ---------- Args ----------
DRY_RUN=0
if [[ "${1:-}" == "--dry-run" ]]; then
  DRY_RUN=1
  shift
fi
(( DRY_RUN == 1 )) && warn "DRY-RUN enabled: no changes will be made."

# shellcheck source=../../lib/universal/niftylib.sh
source <(curl -fSs https://nifty.sh/lib/universal/niftylib.sh)
if ! declare -F require_root &> /dev/null; then
  err2 "Failed to source niftylib"
  exit 1
fi

# ---------- Requirements & policy ----------
require_root
ERR="${ERR:-}"  # guard niftylib internals under set -u
require_cmd awk grep sed tr cut wc df mount blockdev growpart \
            pvs vgs lvs pvdisplay vgdisplay lvdisplay pvresize lvresize \
            resize2fs xfs_growfs

HOST_RE='^(pro)?linux[0-9]+\.(dandomain|wannafind|curanet|scannet|nixfactory)\.dk$'
if printf "%s" "${HOSTNAME:-}" | grep -Pq "$HOST_RE"; then
  ok "Hostname accepted: ${HOSTNAME:-<unknown>}"
else
  err "Hostname '${HOSTNAME:-<unknown>}' not allowed by policy."
  exit 1
fi

# ---------- Policy: minimum extents to add (tune as needed) ----------
# With a typical PE size of 4 MiB, 256 PE ≈ 1 GiB.
MIN_ADD_PE=${MIN_ADD_PE:-256}

# ---------- Helpers ----------
bytes_to_gib() { awk -v b="$1" 'BEGIN {printf "%.3f", b/1024/1024/1024}'; }
bytes_to_mib() { awk -v b="$1" 'BEGIN {printf "%.1f", b/1024/1024}'; }
kb_to_gib()    { awk -v k="$1" 'BEGIN {printf "%.1f", k/1024/1024}'; }

mountpoint_fs()   { awk -v m="$1" '$2==m {print $3}' /proc/mounts | tail -n1; }
mounted_device()  { awk -v m="$1" '$2==m {print $1}' /proc/mounts | tail -n1; }
lv_vg_from_mapper(){ lvs --noheadings -o vg_name,lv_name "$1" 2>/dev/null | awk '{$1=$1;print}'; }
vg_free_extents() { vgs --noheadings -o vg_free_count "$1" 2>/dev/null | awk '{$1=$1;print}'; }
vg_extent_size_b(){ vgs --noheadings --units b --nosuffix -o vg_extent_size "$1" 2>/dev/null | tr -dc '0-9'; }
lv_size_bytes()   { lvs --noheadings --units b --nosuffix -o lv_size "$1" 2>/dev/null | tr -dc '0-9'; }
fs_size_k()       { df -kP "$1" | awk 'NR==2{print $2}'; }
pv_free_bytes()   { pvs --noheadings --units b --nosuffix -o pv_free "$1" 2>/dev/null | tr -dc '0-9'; }

# In dry-run: print the command. In live: run it with stdout/stderr suppressed.
run_or_echo_quiet() {
  if (( DRY_RUN == 1 )); then
    echo "+ $*"
    return 0
  else
    "$@" >/dev/null 2>&1
  fi
}

# Compute how many additional full extents could be added to a PV if resized.
# args: <pv_path> <vg_name>
additional_extents() {
  local pv="$1" vg="$2"
  local pv_b dev_b pe_start_b pe_count pe_sz_b

  read -r pv_b dev_b pe_start_b pe_count < <(
    pvs --noheadings --units b --nosuffix -o pv_size,dev_size,pe_start,pe_count "$pv" \
      2>/dev/null | awk '{print $1, $2, $3, $4}'
  ) || { echo 0; return; }

  pe_sz_b="$(vg_extent_size_b "$vg")"
  if [[ -z "$pv_b" || -z "$dev_b" || -z "$pe_start_b" || -z "$pe_count" || -z "$pe_sz_b" || "$pe_sz_b" = 0 ]]; then
    echo 0; return
  fi

  local new_pe=$(( (dev_b - pe_start_b) / pe_sz_b ))
  local add_pe=$(( new_pe - pe_count ))
  (( add_pe < 0 )) && add_pe=0
  echo "$add_pe"
}

# Return predicted new *bytes* for a partition according to growpart -N.
#  echo value>0  => partition WOULD grow to this size (bytes)
#  echo 0        => no change
#  return 0/1/2  => 0 success (value printed), 1 = no change, 2 = error
growpart_new_size_bytes() {
  local disk="$1" part="$2" out rc sectors new_sz_b ss
  out="$(growpart -N "$disk" "$part" 2>&1)"; rc=$?
  if (( rc != 0 )); then
    warn "growpart -N returned non-zero for $disk $part: $rc"
    echo 0
    return 2
  fi
  if grep -qi 'NOCHANGE' <<<"$out"; then
    echo 0
    return 1
  fi
  # Extract "new: size=NNN" (sectors)
  sectors="$(grep -oE 'new:\s*size=[0-9]+' <<<"$out" | tail -n1 | tr -cd '0-9')"
  if [[ -z "$sectors" ]]; then
    echo 0
    return 1
  fi
  ss="$(blockdev --getss "$disk" 2>/dev/null || echo 512)"
  new_sz_b=$(( sectors * ss ))
  echo "$new_sz_b"
  return 0
}

# Like additional_extents(), but allows overriding dev_size with a predicted value
# args: <pv_path> <vg_name> <dev_bytes_override>
additional_extents_with_devbytes() {
  local pv="$1" vg="$2" dev_override="$3"
  local pv_b pe_start_b pe_count pe_sz_b dev_b new_pe add_pe

  read -r pv_b pe_start_b pe_count < <(
    pvs --noheadings --units b --nosuffix -o pv_size,pe_start,pe_count "$pv" \
      2>/dev/null | awk '{print $1, $2, $3}'
  ) || { echo 0; return; }

  pe_sz_b="$(vg_extent_size_b "$vg")"
  [[ -z "$pe_sz_b" || "$pe_sz_b" = 0 ]] && { echo 0; return; }

  if [[ -n "$dev_override" && "$dev_override" != 0 ]]; then
    dev_b="$dev_override"
  else
    dev_b="$(pvs --noheadings --units b --nosuffix -o dev_size "$pv" 2>/dev/null | tr -dc '0-9')"
  fi

  [[ -z "$dev_b" || -z "$pv_b" || -z "$pe_start_b" || -z "$pe_count" ]] && { echo 0; return; }

  new_pe=$(( (dev_b - pe_start_b) / pe_sz_b ))
  add_pe=$(( new_pe - pe_count ))
  (( add_pe < 0 )) && add_pe=0
  echo "$add_pe"
}

# Return codes for growpart_would_change:
#   0 => partition WOULD grow
#   1 => NO change
#   2 => error running growpart -N
growpart_would_change() {
  local disk="$1" part="$2" out rc
  out="$(growpart -N "$disk" "$part" 2>&1)"; rc=$?
  if (( rc != 0 )); then
    warn "growpart -N returned non-zero for $disk $part: $rc"
    return 2
  fi
  if grep -qi 'NOCHANGE' <<<"$out"; then
    return 1
  fi
  if grep -qiE '(^|\s)CHANGE(D)?\b' <<<"$out" || grep -qi 'old:\s*size=.*\s*new:\s*size=' <<<"$out"; then
    return 0
  fi
  return 1
}

scsi_rescan_all() {
  log "Rescanning SCSI hosts..."
  local any=0
  for h in /sys/bus/scsi/devices/*:*:*:*; do
    if [[ -w "$h/rescan" ]]; then
      printf "%s" "- - -" > "$h/rescan" || true
      any=1
    fi
  done
  if (( any == 1 )); then ok "SCSI rescan requested."; else warn "No SCSI hosts found to rescan."; fi
}

# ---------- Expected layout (targets) ----------
ROOT_MNT="/"
ROOT_LVM_LV="/dev/mapper/vg-root"   # LVM mode
ROOT_VG="vg"
ROOT_DISK="/dev/sda"
ROOT_PV_LVM="/dev/sda2"
ROOT_PLAIN_PART="/dev/sda4"         # non-LVM mode

WWW_MNT="/var/www"
WWW_LV="/dev/mapper/datavg-lv0"
WWW_VG="datavg"
WWW_FS="ext4"
WWW_PV_DISK="/dev/sdc"              # whole disk PV

# ---------- Detect root mode (LVM or plain) ----------
ROOT_DEV="$(mounted_device "$ROOT_MNT")"
ROOT_FS="$(mountpoint_fs "$ROOT_MNT")"
ROOT_MODE=""
if [[ "$ROOT_DEV" == "$ROOT_LVM_LV" ]]; then
  ROOT_MODE="lvm"
  ok "Root detected as LVM: $ROOT_DEV (VG=$ROOT_VG, PV=$ROOT_PV_LVM)"
elif [[ "$ROOT_DEV" == "$ROOT_PLAIN_PART" ]]; then
  ROOT_MODE="plain"
  ok "Root detected as non-LVM: $ROOT_DEV (FS=$ROOT_FS)"
else
  warn "Root mapping $ROOT_DEV not recognized as LVM ($ROOT_LVM_LV) or plain ($ROOT_PLAIN_PART). Root will be skipped."
fi

# ---------- Validate /var/www is as expected ----------
WWW_OK=0
if mountpoint -q "$WWW_MNT"; then
  WWW_DEV="$(mounted_device "$WWW_MNT")"
  WWW_FST="$(mountpoint_fs "$WWW_MNT")"
  if [[ "$WWW_DEV" == "$WWW_LV" && "$WWW_FST" == "$WWW_FS" ]]; then
    if lvs "$WWW_LV" >/dev/null 2>&1; then
      read -r www_vg www_lvname <<<"$(lv_vg_from_mapper "$WWW_LV")"
      if [[ "$www_vg" == "$WWW_VG" && -n "$www_lvname" ]]; then
        if pvs "$WWW_PV_DISK" >/dev/null 2>&1 && pvdisplay "$WWW_PV_DISK" | awk '/VG Name/ {print $3}' | grep -qx "$WWW_VG"; then
          WWW_OK=1
          ok "/var/www validated: $WWW_LV ($WWW_FS) in VG $WWW_VG via PV $WWW_PV_DISK"
        else
          err2 "PV $WWW_PV_DISK not in VG $WWW_VG — skipping /var/www."
        fi
      else
        err2 "$WWW_LV not in expected VG $WWW_VG — skipping /var/www."
      fi
    else
      err2 "LV $WWW_LV not found — skipping /var/www."
    fi
  else
    warn "/var/www mapping unexpected: dev=$WWW_DEV fs=$WWW_FST (expected $WWW_LV/$WWW_FS) — skipping."
  fi
else
  warn "/var/www not mounted — skipping."
fi

# ---------- One global rescan ----------
scsi_rescan_all

# =====================================================================
# Layered “would anything change?” checks with PE math
# =====================================================================

# ---------- /var/www (LVM on /dev/sdc) ----------
process_www() {
  (( WWW_OK == 1 )) || { warn "Skipping /var/www due to pre-checks."; return; }

  # Would we gain at least one full PE on this PV?
  read -r pv_b dev_b pe_start_b pe_count < <(
    pvs --noheadings --units b --nosuffix -o pv_size,dev_size,pe_start,pe_count "$WWW_PV_DISK" 2>/dev/null | awk '{print $1, $2, $3, $4}'
  )
  if [[ -z "${pv_b:-}" || -z "${dev_b:-}" || -z "${pe_start_b:-}" || -z "${pe_count:-}" ]]; then
    warn "Could not read PV/dev size for $WWW_PV_DISK — skipping /var/www."
    return
  fi
  pe_sz_b="$(vg_extent_size_b "$WWW_VG")"
  if [[ -z "${pe_sz_b:-}" || "$pe_sz_b" = 0 ]]; then
    warn "Could not determine VG extent size for $WWW_VG — skipping /var/www."
    return
  fi
  new_pe=$(( (dev_b - pe_start_b) / pe_sz_b ))
  add_pe=$(( new_pe - pe_count ))
  (( add_pe < 0 )) && add_pe=0
  delta_b=$(( add_pe * pe_sz_b ))

  if (( add_pe <= 0 )); then
    ok "/var/www: no full extents available to add (PE=$(bytes_to_mib "$pe_sz_b") MiB). Nothing to do."
    return
  fi
  if (( add_pe < MIN_ADD_PE )); then
    ok "/var/www: would gain +${add_pe} PE (~$(bytes_to_mib "$delta_b") MiB) < MIN_ADD_PE=${MIN_ADD_PE} — skipping by policy."
    return
  fi
  ok "/var/www: would gain +${add_pe} PE (~$(bytes_to_mib "$delta_b") MiB)."

  # Would pvresize actually change metadata?
  if ! pvresize -t "$WWW_PV_DISK" >/dev/null 2>&1; then
    warn "pvresize -t indicates no change (or error) for $WWW_PV_DISK — skipping /var/www."
    return
  fi
  ok "pvresize -t suggests metadata would change."

  # ================= Execute =================
  if ! run_or_echo_quiet pvresize "$WWW_PV_DISK"; then
    err2 "pvresize failed — leaving /var/www unchanged."
    return
  fi
  ok "pvresize completed on $WWW_PV_DISK."

  if (( DRY_RUN == 1 )); then
    run_or_echo_quiet lvresize -l +100%FREE -r "$WWW_LV" --nofsck --yes
    echo "# (Post-verify omitted in DRY-RUN)"
    return
  fi

  # Live: verify VG free and extend LV + FS
  extents="$(vg_free_extents "$WWW_VG")"
  extent_b="$(vg_extent_size_b "$WWW_VG")"; [[ -z "${extent_b:-}" ]] && extent_b=0
  vg_free_b=$(( extents * extent_b ))
  if (( vg_free_b <= 0 )); then
    warn "No free extents in VG $WWW_VG after pvresize — nothing to allocate."
    return
  fi

  if run_or_echo_quiet lvresize -l +100%FREE -r "$WWW_LV" --nofsck --yes; then
    sleep 1
    fs_after_k="$(fs_size_k "$WWW_MNT")"
    ok "/var/www extended. New FS size: $(kb_to_gib "$fs_after_k") GiB"
  else
    err2 "lvresize failed on $WWW_LV — /var/www unchanged."
  fi
}

# ---------- Root (LVM on /dev/sda2) ----------
process_root_lvm() {
  [[ "$ROOT_MODE" == "lvm" ]] || return

  # Would growpart actually enlarge partition 2?
  if growpart_would_change "$ROOT_DISK" 2; then
    ok "growpart -N suggests partition 2 will grow."
  else
    ok "growpart -N reports no partition growth needed for $ROOT_DISK 2 — skipping root (LVM)."
    return
  fi

  # Would pvresize change metadata?
  if ! pvresize -t "$ROOT_PV_LVM" >/dev/null 2>&1; then
    warn "pvresize -t indicates no change for $ROOT_PV_LVM — skipping root (LVM)."
    return
  fi
  ok "pvresize -t suggests PV would grow."

  # Compute predicted new PV extents using growpart -N's *new size* (bytes)
  pred_bytes="$(growpart_new_size_bytes "$ROOT_DISK" 2)"
  pe_sz_b="$(vg_extent_size_b "$ROOT_VG")"

  if [[ -z "$pred_bytes" || "$pred_bytes" = 0 ]]; then
    ok "Root (LVM): growpart predicts no size increase — nothing to do."
    return
  fi

  add_pe="$(additional_extents_with_devbytes "$ROOT_PV_LVM" "$ROOT_VG" "$pred_bytes")"
  delta_b=$(( add_pe * pe_sz_b ))

  if (( add_pe <= 0 )); then
    ok "Root (LVM): predicted increase yields no full PE (PE=$(bytes_to_mib "$pe_sz_b") MiB) — nothing to do."
    return
  fi
  if (( add_pe < MIN_ADD_PE )); then
    ok "Root (LVM): predicted +${add_pe} PE (~$(bytes_to_mib "$delta_b") MiB) < MIN_ADD_PE=${MIN_ADD_PE} — skipping by policy."
    return
  fi
  ok "Root (LVM): would gain +${add_pe} PE (~$(bytes_to_mib "$delta_b") MiB)."

  # ================= Execute =================
  if ! run_or_echo_quiet growpart "$ROOT_DISK" 2; then
    err2 "growpart failed — cannot proceed with root (LVM)."
    return
  fi
  ok "growpart expanded $ROOT_PV_LVM."

  if ! run_or_echo_quiet pvresize "$ROOT_PV_LVM"; then
    err2 "pvresize failed — skipping LV/FS extension for root (LVM)."
    return
  fi
  ok "pvresize completed on $ROOT_PV_LVM."

  if (( DRY_RUN == 1 )); then
    run_or_echo_quiet lvresize -l +100%FREE "$ROOT_LVM_LV" --yes
    run_or_echo_quiet xfs_growfs "$ROOT_MNT"
    echo "# (Post-verify omitted in DRY-RUN)"
    return
  fi

  # Live: verify VG free, extend LV, then grow XFS
  extents="$(vg_free_extents "$ROOT_VG")"
  extent_b="$(vg_extent_size_b "$ROOT_VG")"; [[ -z "${extent_b:-}" ]] && extent_b=0
  vg_free_b=$(( extents * extent_b ))
  if (( vg_free_b <= 0 )); then
    warn "No free extents in VG $ROOT_VG after pvresize — nothing to allocate."
    return
  fi

  if run_or_echo_quiet lvresize -l +100%FREE "$ROOT_LVM_LV" --yes; then
    if run_or_echo_quiet xfs_growfs "$ROOT_MNT"; then
      fs_after_k="$(fs_size_k "$ROOT_MNT")"
      ok "Root (LVM) extended. New FS size: $(kb_to_gib "$fs_after_k") GiB"
    else
      err2 "xfs_growfs failed on / — filesystem unchanged."
    fi
  else
    err2 "lvresize failed — root unchanged."
  fi
}

# ---------- Root (plain /dev/sda4) ----------
process_root_plain() {
  [[ "$ROOT_MODE" == "plain" ]] || return
  local fs="$ROOT_FS"

  # Would growpart actually enlarge partition 4?
  if growpart_would_change "$ROOT_DISK" 4; then
    ok "growpart -N suggests partition 4 will grow."
  else
    ok "growpart -N reports no partition growth needed for $ROOT_DISK 4 — skipping root (plain)."
    return
  fi

  # ================= Execute =================
  if ! run_or_echo_quiet growpart "$ROOT_DISK" 4; then
    err2 "growpart failed for $ROOT_DISK 4 — cannot proceed with root (plain)."
    return
  fi
  ok "growpart expanded $ROOT_PLAIN_PART."

  if (( DRY_RUN == 1 )); then
    if [[ "$fs" == "xfs" ]]; then
      run_or_echo_quiet xfs_growfs "$ROOT_MNT"
    elif [[ "$fs" == "ext4" || "$fs" == "ext3" ]]; then
      run_or_echo_quiet resize2fs "$ROOT_PLAIN_PART"
    else
      err2 "Unsupported root filesystem for plain mode: $fs"
    fi
    echo "# (Post-verify omitted in DRY-RUN)"
    return
  fi

  if [[ "$fs" == "xfs" ]]; then
    if run_or_echo_quiet xfs_growfs "$ROOT_MNT"; then
      fs_after_b="$(df -B1 --output=size "$ROOT_MNT" | awk 'NR==2{print $1}')"
      ok "Root (plain xfs) filesystem grown to $(bytes_to_gib "$fs_after_b") GiB"
    else
      err2 "xfs_growfs failed on /."
    fi
  elif [[ "$fs" == "ext4" || "$fs" == "ext3" ]]; then
    if run_or_echo_quiet resize2fs "$ROOT_PLAIN_PART"; then
      fs_after_b="$(df -B1 --output=size "$ROOT_MNT" | awk 'NR==2{print $1}')"
      ok "Root (plain ext) filesystem grown to $(bytes_to_gib "$fs_after_b") GiB"
    else
      err2 "resize2fs failed on $ROOT_PLAIN_PART."
    fi
  else
    err2 "Unsupported root filesystem for plain mode: $fs"
  fi
}

# ---------- Run targets ----------
process_www
process_root_lvm
process_root_plain

printf "\n"; ok "Processing complete."