#!/bin/bash

### BEGIN INIT INFO
# Provides:          resize2fs
# Required-Start:    $local_fs
# Required-Stop:
# Should-Start:
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: Resize the root filesystem
# Description:	     Resize the root filesystem to fill the whole storage
### END INIT INFO

. /lib/init/vars.sh
. /lib/lsb/init-functions

export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
Log=/var/log/armhwinfo.log

if [ -f /etc/default/FIRSTBOOT ]; then
	# Resize EMMC partitions
	echo "Firstboot: resize rootfs."
	rootpart=$(findmnt -n -o SOURCE /) # i.e. /dev/mmcblk0p1
	resize2fs $rootpart

	rm -f /etc/default/FIRSTBOOT
	# Cleanup firstboot mark after first boot to system
	fw_setenv firstboot 0
	# Delete lock
	rm -rf /var/lock/fw_printenv.lock

	systemctl disable resize2fs
	exit 0
fi

do_expand_partition()
{
	local rootpart=$(findmnt -n -o SOURCE /) # i.e. /dev/mmcblk0p1
	local rootpartname=${rootpart##*/}
	local rootdevice=$(lsblk -n -o PKNAME $rootpart) # i.e. mmcblk0
	local rootdevicepath="/dev/$rootdevice" # i.e. /dev/mmcblk0
	local rootpartsys="/sys/class/block/$rootpartname"
	# get count of partitions and their boundaries without repeatedly
	# invoking parted, which is slow on Ubuntu 26.04 first boot
	local partitions=$(lsblk -n -o PARTN $rootdevicepath | tail -1)
	local partstart=$(cat "$rootpartsys/start")
	local partsize=$(cat "$rootpartsys/size")
	local partend=$((partstart + partsize - 1))
	local startfrom=$(( $partend + 1 ))
	[[ $partitions == 1 ]] && startfrom=$partstart
	local diskbytes=$(blockdev --getsize64 $rootdevicepath)
	local disksectors=$(blockdev --getsz $rootdevicepath)
	local capacity=$(( diskbytes / 1024 / 1024 / 1024 )) # GiB

	# check whether a resizing rule is defined. We will take this value if it's not too low. In
	# this case the value will be ignored and resizing to the whole card size happens.
	if [[ -f /root/.rootfs_resize ]]; then
		read RESIZE_VALUE </root/.rootfs_resize
		ResizeLog="Resize rule $RESIZE_VALUE defined for root partition"
		case $RESIZE_VALUE in
			*%)
				# percentage value, we try to use 16MiB to align partitions since this is
				# the erase block size of more recent SD cards (512 byte sectors, so we use 32
				# as divider and substract 1)
				local percentage=$(echo $RESIZE_VALUE | tr -c -d '[:digit:]')
				local lastsector=$(( 32 * ( disksectors * percentage / 3200 ) - 1 ))
				[[ $lastsector -lt $partend ]] && unset lastsector
				;;
			*s)
				# sector value, we use it directly
				local lastsector=$(echo $RESIZE_VALUE | tr -c -d '[:digit:]')
				[[ $lastsector -lt $partend ]] && unset lastsector
				;;
		esac
		# if SD card is larger than 4GB then create another partition behind first one(s)
		if [[ $capacity -ge 5 ]]; then
			local secondpartition=$(( 32 * ( disksectors * 99 / 3200 ) - 1 ))
			if [[ $secondpartition -lt $partend ]]; then
				unset secondpartition
			fi
		fi
	else
		# check device capacity. If 4GB or below do not use whole card but leave a 5% spare area
		# to help older cards with wear leveling and garbage collection. In case this reduced card
		# capacity is less than the actual image capacity this is a clear sign that someone wants
		# to use Armbian on a card of inappropriate size so he gets what he deserves (at least he
		# should know what he's doing)
		if [[ $capacity -lt 5 ]]; then # 4 GiB or less
			local lastsector=$(( disksectors - 200 * 1024 * 1024 / 512 ))
			if [[ $lastsector -lt $partend ]]; then
				unset lastsector
			else
				ResizeLog="4GB media so leaving 200MB spare area"
			fi
		elif [[ $capacity -lt 9 ]]; then # 8 GiB or less
			# Leave 2 percent unpartitioned
			local lastsector=$(( 32 * ( disksectors * 98 / 3200 ) - 1 ))
			if [[ $lastsector -lt $partend ]]; then
				unset lastsector
			else
				ResizeLog="8GB media so leaving 2 percent spare area"
			fi
		else
			# Leave 1 percent unpartitioned
			local lastsector=$(( 32 * ( disksectors * 99 / 3200 ) - 1 ))
			if [[ $lastsector -lt $partend ]]; then
				unset lastsector
			else
				ResizeLog="Leaving 1 percent spare area"
			fi
		fi
	fi

	# Start resizing
	echo -e "\n### [firstrun] ${ResizeLog}. Start resizing partition $rootpart now:\n" >>${Log}
	cat /proc/partitions >>${Log}
	echo -e "\nExecuting partition resize and partition table reread:" >>${Log}
	if [[ -n $lastsector && -z $secondpartition ]] && command -v sfdisk >/dev/null 2>&1; then
		local newsize=$(( lastsector - partstart + 1 ))
		echo ",$newsize" | sfdisk --force -N $partitions $rootdevicepath >>${Log} 2>&1
	else
		local fdisk_version=$(fdisk --version | awk '{print $NF}' | grep -oE "^[[:digit:]]\.[[:digit:]]+")
		if [[ $partitions == 1 ]] && awk "BEGIN{exit ! ($fdisk_version >= 2.27 )}"; then
			# if dealing with fdisk from util-linux 2.27+ we need a workaround for just 1 partition
			# though it does not break anything - just prevents an "unexpected command" to fdisk
			# https://github.com/igorpecovnik/lib/issues/353#issuecomment-224728506
			((echo d; echo n; echo p; echo ; echo $startfrom; echo $lastsector ; echo w;) | fdisk $rootdevicepath) >>${Log} 2>&1
		else
			((echo d; echo $partitions; echo n; echo p; echo ; echo $startfrom; echo $lastsector ; echo w;) | fdisk $rootdevicepath) >>${Log} 2>&1
		fi
		[[ -n $secondpartition ]] && \
			((echo n; echo p; echo ; echo $(( $lastsector + 1 )); echo $secondpartition ; echo w;) | fdisk $rootdevicepath) >>${Log} 2>&1
	fi
	local s=0
	partx -u $rootdevicepath >>${Log} 2>&1 || partprobe $rootdevicepath >>${Log} 2>&1 || s=$?
	echo -e "\nNew partition table:\n" >>${Log}
	cat /proc/partitions >>${Log}
	echo -e "\nNow trying to resize $1 filesystem on $rootpart to the limits:\n" >>${Log}
	case $1 in
		ext4)
			resize2fs $rootpart >>${Log} 2>&1
			# check whether reboot is necessary for resize2fs to take effect
			local freesize=$(( $(findmnt --target / -n -o AVAIL -b) / 1048576 )) # MiB
			if [[ $s != 0 || $freesize -lt 512 ]]; then
				touch /var/run/resize2fs-reboot
				echo -e "\n### [firstrun] Automated reboot needed to finish the resize procedure" >>${Log}
			fi
			;;
		btrfs)
			btrfs filesystem resize max / >> ${Log} 2>&1
			;;
	esac
}

do_expand_ext4()
{
	echo -e "\n### [resize2fs] Start resizing ext4 partition $1 now\n" >> ${Log}
	resize2fs ${rootpart} >> ${Log} 2>&1
}

do_expand_btrfs()
{
	echo -e "\n### [btrfs resize] Start resizing btrfs partition $1 now\n" >> ${Log}
	btrfs filesystem resize max / >> ${Log} 2>&1
}

case "$1" in
	start)
		# skip resizing if rootfs is neither ext4 not btrfs or if explicitly disabled
		if [[ -f /root/.no_rootfs_resize ]]; then
			systemctl disable resize2fs
			exit 0
		fi

		rootfstype=$(findmnt -n -o FSTYPE /)
		rootpart=$(findmnt -n -o SOURCE /) # i.e. /dev/mmcblk0p1
		case ${rootfstype} in
			ext4)
				# first stage - resize the rootfs partition
				[[ ! -f /var/lib/armbian/resize_second_stage ]] && do_expand_partition ${rootfstype}
				# second stage - resize the filesystem
				[[ ! -f /var/run/resize2fs-reboot ]] && do_expand_ext4 ${rootpart}
				;;
			btrfs)
				do_expand_partition ${rootfstype} && systemctl disable resize2fs
				;;
		esac

		# disable itself
		[[ ! -f /var/run/resize2fs-reboot ]] && systemctl disable resize2fs

		# stop lightdm first run
#		systemctl stop lightdm

		exit 0
		;;

	*)
		echo "Usage: $0 start"
		exit 0
	;;
esac
