#!/usr/bin/env bash
#
# expand_disk
#
# Grow an LVM logical volume by mount point.
#
# Workflow:
#   1. Validate prerequisites
#   2. Ensure growpart is installed
#   3. Resolve mount point -> LV
#   4. Find VG and PV(s)
#   5. Grow backing partition(s) with growpart
#   6. Run pvresize
#   7. Re-read VG free space
#   8. Calculate requested increment
#   9. Grow LV and filesystem
#
# Usage:
#   expand_disk <mountpoint> [increment]
#
# Examples:
#   expand_disk /
#   expand_disk /var/log 25%
#   expand_disk /var 10G
#   expand_disk /home 500M
#   expand_disk /srv 1073741824B
#
# If increment is omitted:
#   100% of remaining VG free space is used.
#
# Percentage values are percentages of VG free space AFTER
# backing partition growth and pvresize.
#

set -euo pipefail

SCRIPT_NAME="$(basename "$0")"

usage() {
    cat <<USAGE
Usage:
  $SCRIPT_NAME <mountpoint> [increment]

Examples:
  $SCRIPT_NAME /
  $SCRIPT_NAME /var/log 25%
  $SCRIPT_NAME /var 10G
  $SCRIPT_NAME /home 500M
  $SCRIPT_NAME /srv 1073741824B

If increment is omitted:
  100% of the remaining VG free space is used.

Supported increments:
  B, K, KB, KBYTE, KBYTES
  M, MB, MBYTE, MBYTES
  G, GB, GBYTE, GBYTES
  percentage, e.g. 25%

Percentage values are calculated after:
  1. backing partition growth
  2. pvresize
  3. updated VG free-space discovery
USAGE
}

die() {
    echo "ERROR: $*" >&2
    exit 1
}

info() {
    echo "==> $*"
}

require_cmd() {
    command -v "$1" >/dev/null 2>&1 ||
        die "Required command not found: $1"
}

trim() {
    local value="$1"

    value="${value#"${value%%[![:space:]]*}"}"
    value="${value%"${value##*[![:space:]]}"}"

    printf '%s' "$value"
}

human_bytes() {
    local bytes="$1"

    awk -v bytes="$bytes" '
    function human(x) {
        if (x >= 1099511627776)
            return sprintf("%.2f TiB", x / 1099511627776)

        if (x >= 1073741824)
            return sprintf("%.2f GiB", x / 1073741824)

        if (x >= 1048576)
            return sprintf("%.2f MiB", x / 1048576)

        if (x >= 1024)
            return sprintf("%.2f KiB", x / 1024)

        return sprintf("%d bytes", x)
    }

    BEGIN {
        print human(bytes)
    }'
}

ensure_growpart() {

    if command -v growpart >/dev/null 2>&1; then
        info "growpart is installed: $(command -v growpart)"
        return 0
    fi

    info "growpart is not installed; attempting to install it"

    if [[ -r /etc/os-release ]]; then
        # shellcheck disable=SC1091
        . /etc/os-release
    else
        die "Unable to determine Linux distribution (/etc/os-release missing)."
    fi

    case "${ID:-}" in

        ubuntu|debian)

            require_cmd apt-get

            info "Detected ${PRETTY_NAME:-$ID}"
            info "Installing cloud-guest-utils"

            apt-get update

            DEBIAN_FRONTEND=noninteractive \
                apt-get install -y cloud-guest-utils
            ;;

        rhel|rocky|almalinux|centos|ol)

            require_cmd dnf

            info "Detected ${PRETTY_NAME:-$ID}"
            info "Installing cloud-utils-growpart"

            dnf install -y cloud-utils-growpart
            ;;

        *)

            if [[ " ${ID_LIKE:-} " == *" debian "* ]]; then

                require_cmd apt-get

                info "Detected Debian-compatible distribution: ${PRETTY_NAME:-${ID:-unknown}}"
                info "Installing cloud-guest-utils"

                apt-get update

                DEBIAN_FRONTEND=noninteractive \
                    apt-get install -y cloud-guest-utils

            elif [[ " ${ID_LIKE:-} " == *" rhel "* ]] ||
                 [[ " ${ID_LIKE:-} " == *" fedora "* ]]; then

                require_cmd dnf

                info "Detected RHEL-compatible distribution: ${PRETTY_NAME:-${ID:-unknown}}"
                info "Installing cloud-utils-growpart"

                dnf install -y cloud-utils-growpart

            else

                die "Unsupported Linux distribution '${PRETTY_NAME:-${ID:-unknown}}'. Install growpart manually."

            fi
            ;;

    esac

    command -v growpart >/dev/null 2>&1 ||
        die "growpart installation completed, but the growpart command is still unavailable."

    info "growpart installed successfully: $(command -v growpart)"
}

# ----------------------------------------------------------------------
# Arguments
# ----------------------------------------------------------------------

if [[ $# -lt 1 || $# -gt 2 ]]; then
    usage
    exit 1
fi

MOUNTPOINT_ARG="$1"
INCREMENT="${2:-100%}"

[[ $EUID -eq 0 ]] ||
    die "This script must be run as root."

MOUNTPOINT_ARG="$(readlink -f "$MOUNTPOINT_ARG")"

mountpoint -q "$MOUNTPOINT_ARG" ||
    die "'$MOUNTPOINT_ARG' is not a mounted filesystem."

# ----------------------------------------------------------------------
# Base dependencies
# ----------------------------------------------------------------------

for cmd in \
    findmnt \
    lsblk \
    lvs \
    vgs \
    pvs \
    pvresize \
    lvextend \
    mountpoint \
    readlink \
    awk \
    sed \
    grep \
    xargs \
    head \
    tr
do
    require_cmd "$cmd"
done

ensure_growpart

# ----------------------------------------------------------------------
# Determine mounted source
# ----------------------------------------------------------------------

SOURCE="$(findmnt -n -o SOURCE --target "$MOUNTPOINT_ARG")"
FSTYPE="$(findmnt -n -o FSTYPE --target "$MOUNTPOINT_ARG")"

[[ -n "$SOURCE" ]] ||
    die "Unable to determine source device for '$MOUNTPOINT_ARG'."

[[ -n "$FSTYPE" ]] ||
    die "Unable to determine filesystem type for '$MOUNTPOINT_ARG'."

SOURCE_REAL="$(readlink -f "$SOURCE")"

[[ -b "$SOURCE_REAL" ]] ||
    die "Mount source '$SOURCE' does not resolve to a block device."

# ----------------------------------------------------------------------
# Locate the LVM LV
#
# Do not decode /dev/mapper names.
#
# Example:
#
#   /dev/mapper/ubuntu--vg-var_log
#
# corresponds to:
#
#   VG = ubuntu-vg
#   LV = var_log
#
# Instead, compare the resolved block devices.
# ----------------------------------------------------------------------

LV_INFO=""

while IFS='|' read -r path vg lv; do

    path="$(trim "$path")"
    vg="$(trim "$vg")"
    lv="$(trim "$lv")"

    [[ -n "$path" ]] || continue

    LV_REAL="$(readlink -f "$path" 2>/dev/null || true)"

    if [[ -n "$LV_REAL" && "$LV_REAL" == "$SOURCE_REAL" ]]; then
        LV_INFO="${path}|${vg}|${lv}"
        break
    fi

done < <(
    lvs \
        --noheadings \
        --separator '|' \
        -o lv_path,vg_name,lv_name
)

[[ -n "$LV_INFO" ]] ||
    die "$SOURCE is not recognized as an LVM logical volume."

IFS='|' read -r LV_PATH VG_NAME LV_NAME <<< "$LV_INFO"

LV_PATH="$(trim "$LV_PATH")"
VG_NAME="$(trim "$VG_NAME")"
LV_NAME="$(trim "$LV_NAME")"

[[ -n "$LV_PATH" ]] ||
    die "Unable to determine logical volume path."

[[ -n "$VG_NAME" ]] ||
    die "Unable to determine volume group."

[[ -n "$LV_NAME" ]] ||
    die "Unable to determine logical volume name."

# ----------------------------------------------------------------------
# Initial status
# ----------------------------------------------------------------------

echo
echo "============================================================"
echo "Examining LVM filesystem"
echo "============================================================"
echo
echo "Mount point       : $MOUNTPOINT_ARG"
echo "Filesystem        : $FSTYPE"
echo "Mounted source    : $SOURCE"
echo "Block device      : $SOURCE_REAL"
echo "Logical volume    : $LV_PATH"
echo "Volume group      : $VG_NAME"
echo "LV name           : $LV_NAME"
echo "Requested increase: $INCREMENT"
echo

# ----------------------------------------------------------------------
# Find PVs belonging to this VG
# ----------------------------------------------------------------------

info "Examining physical volumes in VG '$VG_NAME'"

mapfile -t PV_LIST < <(
    pvs \
        --noheadings \
        --separator '|' \
        -o pv_name,vg_name |
    awk -F'|' -v target="$VG_NAME" '
    {
        pv=$1
        vg=$2

        gsub(/^[ \t]+|[ \t]+$/, "", pv)
        gsub(/^[ \t]+|[ \t]+$/, "", vg)

        if (vg == target)
            print pv
    }'
)

[[ ${#PV_LIST[@]} -gt 0 ]] ||
    die "No physical volumes found for VG '$VG_NAME'."

# ----------------------------------------------------------------------
# Expand backing storage
# ----------------------------------------------------------------------

for PV in "${PV_LIST[@]}"; do

    PV="$(trim "$PV")"

    [[ -n "$PV" ]] || continue

    echo
    echo "------------------------------------------------------------"
    echo "Physical volume: $PV"
    echo "------------------------------------------------------------"

    PV_REAL="$(readlink -f "$PV" 2>/dev/null || true)"

    [[ -n "$PV_REAL" ]] ||
        die "Unable to resolve physical volume '$PV'."

    PV_TYPE="$(lsblk -ndo TYPE "$PV_REAL" 2>/dev/null || true)"
    PV_TYPE="$(trim "$PV_TYPE")"

    echo "Device type      : ${PV_TYPE:-unknown}"

    case "$PV_TYPE" in

        part)

            PARENT_KNAME="$(
                lsblk -ndo PKNAME "$PV_REAL" 2>/dev/null |
                head -n1
            )"

            PART_NUMBER="$(
                lsblk -ndo PARTN "$PV_REAL" 2>/dev/null |
                head -n1
            )"

            PARENT_KNAME="$(trim "$PARENT_KNAME")"
            PART_NUMBER="$(trim "$PART_NUMBER")"

            [[ -n "$PARENT_KNAME" ]] ||
                die "Unable to determine parent disk for $PV."

            [[ -n "$PART_NUMBER" ]] ||
                die "Unable to determine partition number for $PV."

            [[ "$PART_NUMBER" =~ ^[0-9]+$ ]] ||
                die "Invalid partition number '$PART_NUMBER' for $PV."

            PARENT_DISK="/dev/$PARENT_KNAME"

            echo "Parent disk      : $PARENT_DISK"
            echo "Partition number : $PART_NUMBER"

            info "Checking whether $PV can consume additional disk space"

            set +e

            GROW_OUTPUT="$(
                growpart "$PARENT_DISK" "$PART_NUMBER" 2>&1
            )"

            GROW_RC=$?

            set -e

            if [[ -n "$GROW_OUTPUT" ]]; then
                echo "$GROW_OUTPUT" |
                    sed 's/^/    /'
            fi

            if [[ $GROW_RC -eq 0 ]]; then

                info "Partition $PV was expanded"

                if command -v partprobe >/dev/null 2>&1; then
                    partprobe "$PARENT_DISK" 2>/dev/null || true
                fi

                if command -v udevadm >/dev/null 2>&1; then
                    udevadm settle 2>/dev/null || true
                fi

            else

                if echo "$GROW_OUTPUT" |
                    grep -qiE \
                    'NOCHANGE|could only be grown by|not enough free space'
                then
                    info "Partition $PV already consumes the available disk space"
                else
                    die "growpart failed for $PV."
                fi

            fi
            ;;

        disk)

            echo "PV occupies an entire disk."
            echo "Partition expansion is not required."
            ;;

        *)

            echo
            echo "PV type '$PV_TYPE' is not a normal disk partition."
            echo "Automatic partition growth will not be attempted."
            echo
            echo "pvresize will still be attempted."
            ;;

    esac

    echo
    info "Resizing LVM physical volume $PV"

    pvresize "$PV"

done

# ----------------------------------------------------------------------
# Capacity calculations happen only AFTER growpart and pvresize.
# ----------------------------------------------------------------------

echo
echo "============================================================"
echo "Updated volume group geometry"
echo "============================================================"
echo

VG_EXTENT_SIZE="$(
    vgs \
        --noheadings \
        --units b \
        --nosuffix \
        -o vg_extent_size \
        "$VG_NAME" |
    xargs |
    awk '{printf "%.0f", $1}'
)"

VG_FREE_EXTENTS="$(
    vgs \
        --noheadings \
        -o vg_free_count \
        "$VG_NAME" |
    xargs
)"

LV_SIZE_BYTES="$(
    lvs \
        --noheadings \
        --units b \
        --nosuffix \
        -o lv_size \
        "$LV_PATH" |
    xargs |
    awk '{printf "%.0f", $1}'
)"

[[ "$VG_EXTENT_SIZE" =~ ^[0-9]+$ ]] ||
    die "Unable to determine VG extent size."

[[ "$VG_FREE_EXTENTS" =~ ^[0-9]+$ ]] ||
    die "Unable to determine VG free extent count."

[[ "$LV_SIZE_BYTES" =~ ^[0-9]+$ ]] ||
    die "Unable to determine LV size."

VG_FREE_BYTES=$((VG_FREE_EXTENTS * VG_EXTENT_SIZE))

echo "VG                 : $VG_NAME"
echo "Extent size        : $VG_EXTENT_SIZE bytes ($(human_bytes "$VG_EXTENT_SIZE"))"
echo "Free extents       : $VG_FREE_EXTENTS"
echo "VG free space      : $VG_FREE_BYTES bytes ($(human_bytes "$VG_FREE_BYTES"))"
echo "Current LV size    : $LV_SIZE_BYTES bytes ($(human_bytes "$LV_SIZE_BYTES"))"
echo

# ----------------------------------------------------------------------
# Nothing available
# ----------------------------------------------------------------------

if (( VG_FREE_EXTENTS <= 0 )); then

    echo "============================================================"
    echo "No expansion required"
    echo "============================================================"
    echo
    echo "The volume group has no remaining free extents."
    echo

    df -hT "$MOUNTPOINT_ARG"

    exit 0

fi

# ----------------------------------------------------------------------
# Parse requested increment
# ----------------------------------------------------------------------

REQUESTED_EXTENTS=0
REQUESTED_BYTES=0

# ----------------------------------------------------------------------
# Percentage
# ----------------------------------------------------------------------

if [[ "$INCREMENT" =~ ^([0-9]+([.][0-9]+)?)%$ ]]; then

    PERCENT="${BASH_REMATCH[1]}"

    awk -v p="$PERCENT" '
    BEGIN {
        if (p <= 0 || p > 100)
            exit 1
    }' ||
        die "Percentage must be greater than 0 and no more than 100."

    REQUESTED_EXTENTS="$(
        awk \
            -v free="$VG_FREE_EXTENTS" \
            -v pct="$PERCENT" '
        BEGIN {
            n = int((free * pct) / 100)

            if (n < 1 && free > 0)
                n = 1

            if (n > free)
                n = free

            printf "%.0f", n
        }'
    )"

    REQUESTED_BYTES=$((REQUESTED_EXTENTS * VG_EXTENT_SIZE))

# ----------------------------------------------------------------------
# Absolute amount
# ----------------------------------------------------------------------

else

    VALUE="$(
        echo "$INCREMENT" |
        tr '[:lower:]' '[:upper:]' |
        tr -d '[:space:]'
    )"

    NUMBER=""
    MULTIPLIER=""

    if [[ "$VALUE" =~ ^([0-9]+)(B|BYTE|BYTES)?$ ]]; then

        NUMBER="${BASH_REMATCH[1]}"
        MULTIPLIER=1

    elif [[ "$VALUE" =~ ^([0-9]+)(K|KB|KBYTE|KBYTES)$ ]]; then

        NUMBER="${BASH_REMATCH[1]}"
        MULTIPLIER=1024

    elif [[ "$VALUE" =~ ^([0-9]+)(M|MB|MBYTE|MBYTES)$ ]]; then

        NUMBER="${BASH_REMATCH[1]}"
        MULTIPLIER=$((1024 * 1024))

    elif [[ "$VALUE" =~ ^([0-9]+)(G|GB|GBYTE|GBYTES)$ ]]; then

        NUMBER="${BASH_REMATCH[1]}"
        MULTIPLIER=$((1024 * 1024 * 1024))

    else

        die "Invalid increment '$INCREMENT'. Use B, K, M, G, or a percentage such as 25%."

    fi

    (( NUMBER > 0 )) ||
        die "Requested increase must be greater than zero."

    REQUESTED_BYTES=$((NUMBER * MULTIPLIER))

    REQUESTED_EXTENTS=$(( \
        (REQUESTED_BYTES + VG_EXTENT_SIZE - 1) \
        / VG_EXTENT_SIZE \
    ))

    if (( REQUESTED_EXTENTS > VG_FREE_EXTENTS )); then

        echo
        echo "Requested increase : $(human_bytes "$REQUESTED_BYTES")"
        echo "Required extents   : $REQUESTED_EXTENTS"
        echo "Available extents  : $VG_FREE_EXTENTS"
        echo "Available space    : $(human_bytes "$VG_FREE_BYTES")"
        echo

        die "Requested increase exceeds available VG free space."

    fi

    REQUESTED_BYTES=$((REQUESTED_EXTENTS * VG_EXTENT_SIZE))

fi

(( REQUESTED_EXTENTS > 0 )) ||
    die "Calculated increase is zero extents."

NEW_SIZE_BYTES=$((LV_SIZE_BYTES + REQUESTED_BYTES))

# ----------------------------------------------------------------------
# Supported filesystems
# ----------------------------------------------------------------------

case "$FSTYPE" in

    ext2|ext3|ext4|xfs)
        ;;

    *)
        die "Automatic filesystem growth is not enabled for filesystem '$FSTYPE'."
        ;;

esac

# ----------------------------------------------------------------------
# Planned operation
# ----------------------------------------------------------------------

echo
echo "============================================================"
echo "Planned expansion"
echo "============================================================"
echo
echo "Mount point       : $MOUNTPOINT_ARG"
echo "Logical volume    : $LV_PATH"
echo "Volume group      : $VG_NAME"
echo "Filesystem        : $FSTYPE"
echo
echo "Current LV size   : $(human_bytes "$LV_SIZE_BYTES")"
echo "VG free space     : $(human_bytes "$VG_FREE_BYTES")"
echo "Requested         : $INCREMENT"
echo "Extents to add    : $REQUESTED_EXTENTS"
echo "Actual space added: $(human_bytes "$REQUESTED_BYTES")"
echo "New LV size       : $(human_bytes "$NEW_SIZE_BYTES")"
echo

# ----------------------------------------------------------------------
# Expand LV and filesystem
# ----------------------------------------------------------------------

info "Extending logical volume and filesystem"

lvextend \
    -r \
    -l "+${REQUESTED_EXTENTS}" \
    "$LV_PATH"

# ----------------------------------------------------------------------
# Final status
# ----------------------------------------------------------------------

echo
echo "============================================================"
echo "Expansion complete"
echo "============================================================"
echo

echo "Filesystem:"
df -hT "$MOUNTPOINT_ARG"

echo
echo "Logical volume:"
lvs \
    -o lv_name,vg_name,lv_size \
    "$LV_PATH"

echo
echo "Volume group:"
vgs \
    -o vg_name,vg_size,vg_free \
    "$VG_NAME"

echo
echo "Physical volumes:"
pvs \
    -o pv_name,vg_name,pv_size,pv_free \
    --select "vg_name=$VG_NAME"

echo
echo "Done."