#!/usr/bin/env bash
#
# expand_disk
#
# Grow an LVM logical volume by mount point.
#
# Modes:
#
#   Increase mode (default):
#       expand_disk <mountpoint> [increment]
#
#   Absolute mode:
#       expand_disk -a <mountpoint> <target-size>
#       expand_disk --absolute <mountpoint> <target-size>
#
# Examples:
#
#   expand_disk /
#       Consume 100% of remaining VG free space.
#
#   expand_disk /var/log 25%
#       Add 25% of remaining VG free space.
#
#   expand_disk /var 10G
#       Add approximately 10 GiB.
#
#   expand_disk -a /var 20G
#       Make /var approximately 20 GiB total.
#
# Notes:
#
#   * Absolute mode does NOT support percentages.
#   * This script never shrinks an LV or filesystem.
#   * Percentage calculations occur AFTER growpart and pvresize.
#   * If growpart is missing, the script installs it automatically
#     on supported Debian/Ubuntu and RHEL-family distributions.
#

set -euo pipefail

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

# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------

usage() {
    cat <<USAGE
Usage:

  Increase mode:
    $SCRIPT_NAME <mountpoint> [increment]

  Absolute mode:
    $SCRIPT_NAME -a <mountpoint> <target-size>
    $SCRIPT_NAME --absolute <mountpoint> <target-size>

Examples:

  $SCRIPT_NAME /
  $SCRIPT_NAME /var/log 25%
  $SCRIPT_NAME /var 10G
  $SCRIPT_NAME /home 500M

  $SCRIPT_NAME -a /var 20G
  $SCRIPT_NAME --absolute /home 50G

Increase mode:

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

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

Absolute mode:

  Sizes mean the desired FINAL LV size.

  Example:
    Current LV = 8G
    $SCRIPT_NAME -a /var 20G

  Result:
    Approximately 12G is added, producing a 20G LV.

  Percentages are not accepted in absolute mode.

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)
    }'
}

parse_size_bytes() {
    local input="$1"
    local value
    local number
    local multiplier

    value="$(
        printf '%s' "$input" |
        tr '[:lower:]' '[:upper:]' |
        tr -d '[:space:]'
    )"

    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))

    elif [[ "$value" =~ ^([0-9]+)(T|TB|TBYTE|TBYTES)$ ]]; then

        number="${BASH_REMATCH[1]}"
        multiplier=$((1024 * 1024 * 1024 * 1024))

    else
        return 1
    fi

    (( number > 0 )) || return 1

    printf '%s\n' "$((number * multiplier))"
}

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)"
}

# ----------------------------------------------------------------------
# Parse arguments
# ----------------------------------------------------------------------

ABSOLUTE_MODE=0

case "${1:-}" in

    -a|--absolute)
        ABSOLUTE_MODE=1
        shift
        ;;

    -h|--help)
        usage
        exit 0
        ;;

esac

if (( ABSOLUTE_MODE )); then

    if [[ $# -ne 2 ]]; then
        usage
        exit 1
    fi

    MOUNTPOINT_ARG="$1"
    SIZE_SPEC="$2"

    if [[ "$SIZE_SPEC" == *"%" ]]; then
        die "Percentages are not supported in absolute mode."
    fi

else

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

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

fi

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

# ----------------------------------------------------------------------
# Dependencies required before path resolution
# ----------------------------------------------------------------------

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

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

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

# ----------------------------------------------------------------------
# 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 LVM LV
#
# Do not decode /dev/mapper escaping.
#
# Instead, compare the resolved block device for each canonical
# LVM LV path against the mounted source.
# ----------------------------------------------------------------------

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"

if (( ABSOLUTE_MODE )); then
    echo "Mode              : absolute"
    echo "Requested size    : $SIZE_SPEC"
else
    echo "Mode              : increase"
    echo "Requested increase: $SIZE_SPEC"
fi

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 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

# ----------------------------------------------------------------------
# Calculate requested change
# ----------------------------------------------------------------------

REQUESTED_EXTENTS=0
REQUESTED_BYTES=0
TARGET_BYTES=0

if (( ABSOLUTE_MODE )); then

    # ------------------------------------------------------------------
    # Absolute mode
    # ------------------------------------------------------------------

    TARGET_BYTES="$(parse_size_bytes "$SIZE_SPEC")" ||
        die "Invalid absolute size '$SIZE_SPEC'. Use B, K, M, G, or T."

    #
    # Never shrink.
    #
    if (( TARGET_BYTES < LV_SIZE_BYTES )); then

        echo "Current LV size : $(human_bytes "$LV_SIZE_BYTES")"
        echo "Requested size  : $(human_bytes "$TARGET_BYTES")"

        die "Requested absolute size is smaller than the current LV. Shrinking is not supported."

    fi

    #
    # Already at or beyond requested target.
    #
    if (( TARGET_BYTES == LV_SIZE_BYTES )); then

        echo "============================================================"
        echo "No expansion required"
        echo "============================================================"
        echo
        echo "Current LV size : $(human_bytes "$LV_SIZE_BYTES")"
        echo "Requested size  : $(human_bytes "$TARGET_BYTES")"
        echo "Action required : NONE"
        echo

        df -hT "$MOUNTPOINT_ARG"

        exit 0

    fi

    #
    # Because LVM operates in extents, the existing LV size should
    # already be extent-aligned.
    #
    # Calculate the required total target extents by rounding the
    # requested target UP to the nearest extent.
    #
    CURRENT_EXTENTS=$((LV_SIZE_BYTES / VG_EXTENT_SIZE))

    TARGET_EXTENTS=$(( \
        (TARGET_BYTES + VG_EXTENT_SIZE - 1) \
        / VG_EXTENT_SIZE \
    ))

    REQUESTED_EXTENTS=$((TARGET_EXTENTS - CURRENT_EXTENTS))

    (( REQUESTED_EXTENTS > 0 )) || {

        echo "============================================================"
        echo "No expansion required"
        echo "============================================================"
        echo
        echo "Requested target falls within the LV's current extent."
        echo "Action required : NONE"
        echo

        exit 0
    }

    if (( REQUESTED_EXTENTS > VG_FREE_EXTENTS )); then

        REQUIRED_BYTES=$((REQUESTED_EXTENTS * VG_EXTENT_SIZE))

        echo
        echo "Current LV size    : $(human_bytes "$LV_SIZE_BYTES")"
        echo "Requested LV size  : $(human_bytes "$TARGET_BYTES")"
        echo "Space required     : $(human_bytes "$REQUIRED_BYTES")"
        echo "VG free space      : $(human_bytes "$VG_FREE_BYTES")"
        echo "Required extents   : $REQUESTED_EXTENTS"
        echo "Available extents  : $VG_FREE_EXTENTS"
        echo

        die "Insufficient VG free space to reach requested absolute size."

    fi

    REQUESTED_BYTES=$((REQUESTED_EXTENTS * VG_EXTENT_SIZE))
    NEW_SIZE_BYTES=$((LV_SIZE_BYTES + REQUESTED_BYTES))

else

    # ------------------------------------------------------------------
    # Increase mode
    # ------------------------------------------------------------------

    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

    #
    # Percentage increase
    #
    if [[ "$SIZE_SPEC" =~ ^([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))

    else

        #
        # Absolute byte increment
        #
        REQUESTED_BYTES="$(parse_size_bytes "$SIZE_SPEC")" ||
            die "Invalid increment '$SIZE_SPEC'. Use B, K, M, G, T, or a percentage such as 25%."

        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

        #
        # Actual allocation is extent aligned.
        #
        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))

fi

# ----------------------------------------------------------------------
# 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")"

if (( ABSOLUTE_MODE )); then

    echo "Mode               : absolute"
    echo "Requested LV size  : $(human_bytes "$TARGET_BYTES")"
    echo "Actual final size  : $(human_bytes "$NEW_SIZE_BYTES")"

else

    echo "Mode               : increase"
    echo "Requested          : $SIZE_SPEC"
    echo "Actual space added : $(human_bytes "$REQUESTED_BYTES")"
    echo "New LV size        : $(human_bytes "$NEW_SIZE_BYTES")"

fi

echo "Extents to add     : $REQUESTED_EXTENTS"
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."