#!/usr/bin/env bash

VERS="0.3.4"
DELAY=5
FREE_SPACE_LIMIT=200

[ -t 1 ] && . colors

set -o pipefail

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

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

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"

	echo -e "${CYAN}> Resizing filesystem${NOCOLOR}"
	resize2fs $dev_part # -f
	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

	if [[ "$(echo $part_name | grep -c nvme)" -gt 0 ]]; then
	    local part_num=`echo $2 | sed -e "s/^$1//" | sed 's/^p//'`
	else
	    local part_num=`echo $2 | sed 's/\([[:alpha:]]*\)//'`
	fi

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

#Just a copy-paste
#bootpart=`readlink -f /dev/block/$(mountpoint -d /)`
#bootdisk=${bootpart::-1} #cutoff partnumber
GPT=0
PART_UUID=`cat /proc/cmdline | tr " " "\n" | grep "UUID" | sed 's/\(^root=UUID=\)//'`
DISK_PART=`blkid -U "${PART_UUID}" | sed 's/\(^\/dev\/\)//'`
if [[ "$(echo $DISK_PART | grep -c nvme)" -gt 0 ]];
then
	DISK_NAME=${DISK_PART::-2} #cutoff partnumber pN
else
	DISK_NAME=${DISK_PART::-1} #cutoff partnumber N
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}" | grep -oP "^Block (count|size):[\s]+\K[0-9]+" | awk -v RS="" '{print int($1*$2/1024/1024+0.5)}'` || 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
CONF_SIZE=$(blockdev --getsize64 $(blkid | grep -E "${DISK_NAME}" | grep -m 1 -E "ntfs" | awk '{print $1}' | sed 's/:$//') | awk '{printf ("%0.f",$1/1024/1024)}')

# Booot (grub) partition, usually 1MB
BOOT_SIZE=$(blockdev --getsize64 $(fdisk -l | grep -E "$DISK_NAME" | grep -E "BIOS" | awk '{print $1}') | awk '{printf ("%0.f",$1/1024/1024)}')

if [[ $GPT == 1 ]]; then
        # UEFI partition, ~40MB on Images prior 20.04 based, ~200MB on 20.04
	UEFI_SIZE=$(blockdev --getsize64 $(blkid | grep -E "${DISK_NAME}" | grep -E "EFI" | awk '{print $1}' | sed 's/:$//') | awk '{printf ("%0.f",$1/1024/1024)}')
	# Linux partition
	EXPAND_SIZE=$((DISK_SIZE-PART_SIZE-CONF_SIZE-UEFI_SIZE-BOOT_SIZE-1)) # minus Hive configuration partition and EFI part and reserved
else
	EXPAND_SIZE=$((DISK_SIZE-PART_SIZE-BOOT_SIZE-CONF_SIZE-1)) # minus Hive configuration partition and reserved
fi

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 for break process in interactive mode
	if [[ -z $SILENT ]]; then
		echo -e "Press Ctrl+C to stop or wait 5 seconds to continue"
		DelayDot
	fi

	install-deps
	part_expand "$DISK_NAME" "$DISK_PART"
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"
else
	[[ -z $SILENT ]] && echo -e "${RED}Free disk space less then ${FREE_SPACE_LIMIT} MB. Exiting${NOCOLOR}"
	exit 1
fi

if [[ $? -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}" | 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"
fi

exit 0
