#!/usr/bin/env bash

VERS="0.3.5"
DELAY=5
FREE_SPACE_LIMIT=200

[ -t 1 ] && . colors

set -o pipefail

[[ $1 == "-s" || $1 == "--silent" ]] && SILENT=1

function DelayDot(){
	local i=1
	while [ "$i" -le "$DELAY" ]
	do
		echo -n '.'
		sleep 1
		i=$((i+1))
	done
	echo
}

# Size in MB of a block device; 0 if the device path is empty or missing.
part_size_mb(){
	local dev=$1 sz
	[[ -z $dev || ! -b $dev ]] && { echo 0; return; }
	sz=$(blockdev --getsize64 "$dev" 2>/dev/null | awk '{printf("%.0f",$1/1024/1024)}')
	echo "${sz:-0}"
}

fs_expand(){
	local part_name=$1
	[[ -z $part_name ]] && echo -e "${YELLOW}No partition name given${NOCOLOR}" && return 1
	local dev_part="/dev/$part_name"
	local fstype
	fstype=$(blkid -o value -s TYPE "$dev_part" 2>/dev/null)

	echo -e "${CYAN}> Resizing filesystem (${fstype:-unknown})${NOCOLOR}"
	case "$fstype" in
		ext2|ext3|ext4)
			resize2fs "$dev_part" ;;
		xfs)
			# xfs grows online only, addressed by mountpoint (OS partition = /)
			xfs_growfs / ;;
		btrfs)
			btrfs filesystem resize max / ;;
		*)
			echo -e "${RED}Unsupported filesystem '${fstype:-unknown}' - cannot resize${NOCOLOR}"
			return 1 ;;
	esac
	if [ $? = 0 ]; then
		echo -e "${GREEN}Resizing complete${NOCOLOR}"
	else
		echo -e "${RED}Resizing was not completed${NOCOLOR}"
		return 1
	fi

	return 0
}

part_expand(){
	[[ -z $1 ]] && echo -e "${YELLOW}No drive name given${NOCOLOR}" && return 1
	local dev_name="/dev/$1"
	local part_name=$2
	[[ -z $part_name ]] && echo -e "${YELLOW}No partition name given${NOCOLOR}" && return 1

	# Partition number straight from sysfs - robust for sd*, nvme, mmcblk and any index
	local part_num
	part_num=$(cat "/sys/class/block/$part_name/partition" 2>/dev/null)
	[[ -z $part_num ]] && echo -e "${RED}Could not determine partition number for $part_name${NOCOLOR}" && return 1

	echo -e "${CYAN}> Starting to repartition${NOCOLOR}"
	growpart $dev_name $part_num
	if [ $? = 0 ]; then
		echo -e "${GREEN}Repartition complete${NOCOLOR}"
	else
		echo -e "${RED}Repartition was not completed${NOCOLOR}"
		return 1
	fi

	fs_expand $part_name
	return
}


#growpart is required for this to work
function install-deps() {
	dpkg -s cloud-guest-utils > /dev/null 2>&1
	[[ $? -ne 0 ]] &&
		apt install -y cloud-guest-utils &&
		echo ""

	which growpart > /dev/null 2>&1
	[[ $? -ne 0 ]] &&
		echo "growpart not found, try to reinstall cloud-guest-utils package" &&
		exit 1
}

######################################################################

[[ -z $SILENT ]] &&
	echo -e "${CYAN}Disk Expand Utility v${VERS}${NOCOLOR}"

if [[ $EUID -ne 0 ]]; then
	echo -e "${RED}This utility must be run as root${NOCOLOR}"
	exit 1
fi

GPT=0
# Root partition device (e.g. sda1, nvme0n1p2, mmcblk0p2). Prefer what is really
# mounted at / - works for both root=UUID= and root=/dev/xxx boot styles; fall
# back to the UUID passed on the kernel cmdline.
DISK_PART=$(findmnt -no SOURCE / 2>/dev/null | sed 's|^/dev/||')
if [[ -z $DISK_PART || ! -b /dev/$DISK_PART ]]; then
	PART_UUID=$(tr " " "\n" < /proc/cmdline | grep -m1 "^root=UUID=" | sed 's/^root=UUID=//')
	[[ -n $PART_UUID ]] && DISK_PART=$(blkid -U "$PART_UUID" | sed 's|^/dev/||')
fi
if [[ -z $DISK_PART || ! -b /dev/$DISK_PART ]]; then
	echo -e "${RED}Could not determine the root partition${NOCOLOR}"
	exit 1
fi

# Parent disk via sysfs/lsblk - robust for sd*, nvme, mmcblk (no char-cutoff guessing)
DISK_NAME=$(lsblk -no pkname "/dev/$DISK_PART" 2>/dev/null | head -n1)
if [[ -z $DISK_NAME || ! -b /dev/$DISK_NAME ]]; then
	echo -e "${RED}Could not determine the parent disk of /dev/$DISK_PART${NOCOLOR}"
	exit 1
fi

[[ $(blkid -o value -s PTTYPE /dev/${DISK_NAME}) == "gpt" ]] && GPT=1
DISK_SIZE=`cat /proc/partitions | grep -w "${DISK_NAME}" | awk '{ printf("%.f",$3/1024)}'`
PART_SIZE=`cat /proc/partitions | grep -w "${DISK_PART}" | awk '{ printf("%.f",$3/1024)}'`
# in some cases filesystem can be smaller than partition size
FS_SIZE=`tune2fs -l "/dev/${DISK_PART}" 2>/dev/null | grep -oP "^Block (count|size):[\s]+\K[0-9]+" | awk -v RS="" '{print int($1*$2/1024/1024+0.5)}'` || FS_SIZE=$PART_SIZE
[[ -z $FS_SIZE ]] && FS_SIZE=$PART_SIZE
FS_TYPE=`df --output=fstype / 2>/dev/null | tail -n 1` || FS_TYPE="unknown"
FS_EXPAND=$(( PART_SIZE - FS_SIZE ))

if [[ -z $SILENT ]]; then
	echo -e "OS booted from ${GREEN}$DISK_NAME${NOCOLOR} and whole drive size is ${WHITE}$DISK_SIZE MB${NOCOLOR}"
	echo -e "OS partition is ${YELLOW}$DISK_PART${NOCOLOR} and it's size is ${WHITE}$PART_SIZE MB${NOCOLOR}"
	echo -e "OS partition filesystem is ${PURPLE}$FS_TYPE${NOCOLOR} and it's size is ${WHITE}$FS_SIZE MB${NOCOLOR}"
fi

# Config partition, usually 20MB (ntfs) - 0 if the image has none
CONF_DEV=$(blkid | grep -E "${DISK_NAME}" | grep -m1 -E "ntfs" | awk '{print $1}' | sed 's/:$//')
CONF_SIZE=$(part_size_mb "$CONF_DEV")

# Boot (grub BIOS) partition, usually 1MB - 0 if absent
BOOT_DEV=$(fdisk -l 2>/dev/null | grep -E "$DISK_NAME" | grep -m1 -E "BIOS" | awk '{print $1}')
BOOT_SIZE=$(part_size_mb "$BOOT_DEV")

if [[ $GPT == 1 ]]; then
	# UEFI/EFI partition (~40MB on pre-20.04 images, ~200MB on 20.04+)
	UEFI_DEV=$(blkid | grep -E "${DISK_NAME}" | grep -m1 -E "EFI" | awk '{print $1}' | sed 's/:$//')
	UEFI_SIZE=$(part_size_mb "$UEFI_DEV")
	# Linux partition = disk - OS part - config - EFI - BIOS - 1MB reserved
	EXPAND_SIZE=$((DISK_SIZE-PART_SIZE-CONF_SIZE-UEFI_SIZE-BOOT_SIZE-1))
else
	EXPAND_SIZE=$((DISK_SIZE-PART_SIZE-BOOT_SIZE-CONF_SIZE-1))
fi

RESULT=0
if [[ $EXPAND_SIZE -gt ${FREE_SPACE_LIMIT} ]]; then
	echo -e "${YELLOW}> Trying to auto-expand OS partition by ${WHITE}${EXPAND_SIZE} MB${NOCOLOR}${YELLOW} up to ${WHITE}$((PART_SIZE+EXPAND_SIZE)) MB${NOCOLOR}${YELLOW} in total${NOCOLOR}"

	# Give user time to break the process in interactive mode
	if [[ -z $SILENT ]]; then
		echo -e "Press Ctrl+C to stop or wait ${DELAY} seconds to continue"
		DelayDot
	fi

	install-deps
	part_expand "$DISK_NAME" "$DISK_PART"
	RESULT=$?
elif [[ $FS_EXPAND -gt ${FREE_SPACE_LIMIT} ]]; then
	echo -e "${YELLOW}> Trying to auto-expand OS filesystem by ${WHITE}$FS_EXPAND MB${NOCOLOR}${YELLOW} up to ${WHITE}$((FS_SIZE+FS_EXPAND)) MB${NOCOLOR}${YELLOW} in total${NOCOLOR}"
	fs_expand "$DISK_PART"
	RESULT=$?
else
	[[ -z $SILENT ]] && echo -e "${RED}Free disk space less then ${FREE_SPACE_LIMIT} MB. Exiting${NOCOLOR}"
	exit 1
fi

if [[ $RESULT -eq 0 ]]; then
	partprobe > /dev/null 2>&1
	NEW_PART_SIZE=`cat /proc/partitions | grep -w "${DISK_PART}" | awk '{ printf("%.f",$3/1024)}'`
	NEW_FS_SIZE=`tune2fs -l "/dev/${DISK_PART}" 2>/dev/null | grep -oP "^Block (count|size):[\s]+\K[0-9]+" | awk -v RS="" '{print int($1*$2/1024/1024+0.5)}'` || NEW_FS_SIZE=$NEW_PART_SIZE
	echo -e "New OS partition size is ${GREEN}$NEW_PART_SIZE${NOCOLOR} MB and filesystem size is ${GREEN}$NEW_FS_SIZE${NOCOLOR} MB"
else
	echo -e "${RED}Disk expand failed${NOCOLOR}"
fi

exit $RESULT
