#!/bin/bash
#
# Copyright (c) Authors: https://www.armbian.com/authors
#
# This file is licensed under the terms of the GNU General Public
# License version 2. This program is licensed "as is" without any
# warranty of any kind, whether express or implied.

. /etc/armbian-release
. /lib/init/vars.sh
. /lib/lsb/init-functions
. /usr/lib/armbian/armbian-common

do_expand_partition()
{
	local partdev=$1
	echo -e "\n### [resize2fs] Trying to resize partition $partdev:\n"

	# make sure the target is a partition
	local parttype=$(lsblk -n -o TYPE $partdev | head -1)
	if [[ "$parttype" != "part" ]]; then
		echo -e "\n$partdev isn't a partition: $parttype" >&2
		return 1
	fi

	local diskdevname=$(lsblk -n -o PKNAME $partdev | head -1) # i.e. mmcblk0 or sda
	# due to the bug in util-linux 2.34 which fails to show device, let's use this failover:
	[[ -z $diskdevname ]] && diskdevname=$(echo $partdev | sed -e "s/^\/dev\///" | sed "s/p.*//")
	local diskdev="/dev/$diskdevname" # i.e. /dev/mmcblk0, /dev/sda
	echo "diskdevname: $diskdevname"
	echo "diskdev: $diskdev"

	# detect partindex
	local partindex="$(echo $partdev | sed "s|^$diskdev||" | sed 's/^p//')" # i.e. 1
	if [[ -n "$(echo "$partindex" | tr -d '[:digit:]')" ]]; then
		echo -e "\nFail to detect partindex: $partindex" >&2
		return 1
	fi
	echo "partindex: $partindex"

	# find the target part's boundaries
	local partstart=$(parted $diskdev unit s print -sm | awk -F':' "NR>=3{if(\$1 == $partindex){print \$2}}" | sed 's/s//')
	local partend=$(parted $diskdev unit s print -sm | awk -F':' "NR>=3{if(\$1 == $partindex){print \$3}}" | sed 's/s//')
	if [[ -z "$partstart" ]] || [[ -z "$partend" ]]; then
		echo -e "\nFail to find the target partition" >&2
		return 1
	fi
	echo "partstart: ${partstart}s"
	echo "partend: ${partend}s"

	# check if it's the last partition in logical layout
	# Note: the order may be different between partition table and logical layout
	local list=$(parted $diskdev unit s print -sm | awk -F':' "NR>=3{if(\$2 - $partstart > 0){print \$1}}")
	if [[ -n "$list" ]]; then
		echo -e "\nThe target partition isn't the last partition in logical layout" >&2
		echo "Partition ${list[*]} behind it" >&2
		return 1
	fi

	local capacity=$(parted $diskdev unit s print -sm | awk -F':' 'NR==2{print $2}' | sed 's/s//')
	local sectorsize=$(parted $diskdev unit s print -sm | awk -F':' 'NR==2{print $4}')
	echo "capacity: ${capacity}s"
	echo "sectorsize: ${sectorsize}B"

	# 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
		echo "Resize rule $RESIZE_VALUE defined for root partition"
		case $RESIZE_VALUE in
			*%)
				# percentage value
				local percentage=$(echo $RESIZE_VALUE | tr -c -d '[:digit:]')
				local lastsector=$(( $capacity * $percentage / 100 ))
				;;
			*s)
				# sector value, we use it directly
				local lastsector=$(echo $RESIZE_VALUE | tr -c -d '[:digit:]')
				;;
			*B)
				# byte value
				local byte=$(echo $RESIZE_VALUE | tr -c -d '[:digit:]')
				local lastsector=$(( $byte / $sectorsize ))
				;;
			*)
				echo -e "\nUnknown rule: $RESIZE_VALUE" >&2
				return 1;
				;;
		esac

		# if remain more than 1GiB then create another partition
		if [[ $(( $sectorsize * ($capacity - $lastsector) )) -ge $(( 1 * 1024 * 1024 * 1024 )) ]]; then
			local newpartition=$(( $capacity * 99 / 100 ))
		fi
	else
		# check device capacity. If 4GiB 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 [[ $(( $sectorsize * $capacity )) -lt $(( 4 * 1024 * 1024 * 1024 )) ]]; then # 4 GiB or less
			# Leave 5 percent unpartitioned
			local lastsector=$(( $capacity * 95 / 100 ))
			echo "4GiB or smaller media - leaving 5% spare area"
		elif [[ $(( $sectorsize * $capacity )) -lt $(( 8 * 1024 * 1024 * 1024 )) ]]; then # 8 GiB or less
			# Leave 2 percent unpartitioned
			local lastsector=$(( $capacity * 98 / 100 ))
			echo "8GiB or smaller media - leaving 2% spare area"
		else
			# Leave 1 percent unpartitioned
			local lastsector=$(( $capacity * 99 / 100 ))
			echo "Leaving 1% spare area"
		fi
	fi

	# use 16MiB to align partitions since this is the erase block size of more recent SD cards
	local lastsector=$(( $lastsector - ($lastsector % (16 * 1024 * 1024 / $sectorsize)) ))
	[[ -n "$newpartition" ]] && local newpartition=$(( $newpartition - ($newpartition % (16 * 1024 * 1024 / $sectorsize)) ))

	if [[ $lastsector -lt $partend ]]; then
		echo -e "\n${partend} is larger than ${lastsector}s, that means you trying to shrink filesystem. Sorry, this tool designed for expanding only." >&2
		return 1
	elif [[ $lastsector -eq $partend ]]; then
		echo -e "\nSame size. Skip" >&2
		return 0
	fi
	if [[ $newpartition -le $(( $lastsector + 1 )) ]]; then
		unset newpartition
	fi

	echo -e "Old partition table:\n"
	cat /proc/partitions

	# Start resizing
	echo -e "\nExecuting fdisk:"
	local partnum=$(parted $diskdev unit s print -sm | awk 'END{print NR-2}')
	local fdisk_version=$(fdisk --version | awk '{print $NF}' | grep -oE "^[[:digit:]]\.[[:digit:]]+")
	if [[ $partnum == 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 $partindex; echo $partstart; echo $lastsector
			echo w
		} | fdisk $diskdev
	else
		{
			echo d; echo $partindex
			echo n; echo p; echo $partindex; echo $partstart; echo $lastsector
			echo w
		} | fdisk $diskdev
	fi
	if [[ -n $newpartition ]]; then
		{
			echo n; echo p; echo ; echo $(( $lastsector + 1 )); echo $newpartition
			echo w
		} | fdisk $diskdev
	fi

	echo -e "\nExecuting partprobe:"
	# Workaround for Kernel bug in 5.8.y and up. Ignore partprobe returning error and inticating that fs is not expended while it is
	KERNELID=$(uname -r |  awk -F'.' '{print ($1 * 100) + $2}')
	if ! partprobe $rootdev && [[ ${KERNELID} -le 507 ]]; then
		echo -e "\n### [resize2fs] Automated reboot needed to finish the resize procedure"
		touch /var/run/resize2fs-reboot
	fi

	local partsize=$(lsblk -n -b -o SIZE $partdev | head -1)
	local actualsize=$(( $sectorsize * ($lastsector - $partstart + 1) ))
	if [[ $actualsize -ne $partsize ]]; then
		echo -e "\n### [resize2fs] Automated reboot needed to finish the resize procedure"
		touch /var/run/resize2fs-reboot
	fi

	echo -e "New partition table:\n"
	cat /proc/partitions

	return 0
}

do_resize_crypt()
{
	local dev=$1
	local parentdev=$2
	local name=$(lsblk -n -o NAME $dev | head -1) # i.e. armbian-root

	echo -e "\n### [resize2fs] Start resizing LUKS container now\n"
	# It's probably no need to run 'cryptsetup resize'.
	# After reboot, it will auto resize to adapte the partition
	# 'cryptsetup resize' requires passphrase, so it will fail.
	cryptsetup resize $name

	local parentsize=$(lsblk -n -b -o SIZE $parentdev | head -1)

	local sectorsize=$(cryptsetup status $name | awk -F':' '/^ *sector size/ {print $2}' | tr -c -d '[:digit:]')
	local offset=$(cryptsetup status $name | awk -F':' '/^ *offset/ {print $2}' | tr -c -d '[:digit:]')
	local size=$(cryptsetup status $name | awk -F':' '/^ *size/ {print $2}' | tr -c -d '[:digit:]')
	local actualsize=$(( $sectorsize * ($offset + $size) ))
	if [[ $actualsize -ne $parentsize ]]; then
		echo -e "\n### [resize2fs] Automated reboot needed to finish the resize procedure"
		touch /var/run/resize2fs-reboot
	fi
}

do_expand_filesystem()
{
	local partdev=$1
	local mountpoint=$2

	local fstype=$(findmnt -n -o FSTYPE $mountpoint) # i.e. ext4 or btrfs

	echo -e "\n### [resize2fs] Trying to resize $fstype filesystem on $partdev:\n"

	case $fstype in
		ext4)
			echo "Running 'resize2fs $partdev' now..."
			resize2fs $partdev
			;;
		btrfs)
			echo "Running 'btrfs filesystem resize max $mountpoint' now..."
			btrfs filesystem resize max $mountpoint
			;;
		nilfs2)
			echo "Running 'nilfs2 filesystem resize max $mountpoint' now..."
			nilfs-resize -v -y $partdev
			;;
		*)
			echo "Unsupported filesystem: $fstype"
			echo "Trying to run 'resize2fs $partdev' now..."
			resize2fs $partdev
			return
			;;
	esac

	# check whether reboot is necessary for resize2fs to take effect
	local devsize=$(lsblk -n -b -o SIZE $partdev | head -1)
	local fssize=$(findmnt -n -b -o SIZE --target $mountpoint)
	if [[ $(( 100 * $fssize / $devsize )) -lt 90 ]]; then
		echo -e "\n### [resize2fs] Automated reboot needed to finish the resize procedure"
		touch /var/run/resize2fs-reboot
	fi
}

main()
{
	rm /var/run/resize2fs-reboot 2>/dev/null

	# skip resizing if explicitly disabled
	if [[ -f /root/.no_rootfs_resize ]]; then
		systemctl disable armbian-resize-filesystem
		exit 0
	fi

	# Logging header
	CPU_ARCH=$(lscpu | awk '/Architecture/ {print $2}')
	DISTRO_ARCH=$(dpkg --print-architecture)
	KERNELID=$(uname -r)
	echo "$(date) | ${BOARD_NAME} | ${VERSION} | ${DISTRO_ARCH} | ${CPU_ARCH} | ${KERNELID}"

	# trim any btrfs subvolume identifier given in square brackets (e.g. /dev/mapper/armbian-root[/@])
	local rootdev=$(findmnt -n -o SOURCE / | sed 's~\[.*\]~~') # i.e. /dev/mmcblk0p1, /dev/sda1 or /dev/mapper/armbian-root

	# check for crypt
	local cryptname=""
	if [[ "$(lsblk -n -o TYPE $rootdev | head -1)" == "crypt" ]]; then
		local cryptname=$(lsblk -n -o NAME $rootdev | head -1) # i.e. armbian-root
		local parent_uuid=$(cat /etc/crypttab | awk "{if(\$1 == \"$cryptname\"){print \$2}}" | sed 's/UUID=//')
		local parentdev=$(blkid -U $parent_uuid) # i.e. /dev/mmcblk0p1 or /dev/sda1

		do_expand_partition "$parentdev"

		# resize LUKS container
		do_resize_crypt "$rootdev" "$parentdev"
	else
		# check if growroot (from cloud-initramfs-growroot package) is installed.
		# despite it's name, that package does NOT require cloud-init.
		# it will resize the *partition* which root filesystem on during initramfs.
		# in this case do nothing here, just run resize2fs (growroot does not handle that).
		# Note: it can't handle crypt device
		if [[ -f /usr/share/initramfs-tools/hooks/growroot ]] || [[ -f /usr/share/initramfs-tools/scripts/local-bottom/growroot ]]; then
			echo "partition resize skipped: growroot detected."
		else
			do_expand_partition "$rootdev"
		fi
	fi

	do_expand_filesystem "$rootdev" "/"

	# disable itself
	[[ ! -f /var/run/resize2fs-reboot ]] && systemctl disable armbian-resize-filesystem
	exit 0
}

case "$1" in
	start)
		touch ${Log}
		chmod 644 ${Log}
		main |& tee -a ${Log}
		;;
	*)
		echo "Usage: $0 start"
		exit 0
	;;
esac