#!/bin/sh

# Uncomment below to see more logs
# set -x

BUSYBOX_MOUNT_OPTS="loop (a|)sync (no|)atime (no|)diratime (no|)relatime (no|)dev (no|)exec (no|)suid (r|)shared (r|)slave (r|)private (un|)bindable (r|)bind move remount ro"
NTFS_3G_MOUNT_OPTS="ro uid=[0-9]* gid=[0-9]* umask=[0-9]* fmask=[0-9]* dmask=[0-9]*"

message()
{
	echo "${ID:+"[$ID]: "}$@"
}

check_tool()
{
	TOOL=${1%% *}
	BR2_CONFIG=$2

	type $TOOL >/dev/null && return 0

	if grep -wq "ID=buildroot" /etc/os-release 2>/dev/null; then
		[ -n "$BR2_CONFIG" ] && \
			message "You may need to enable $BR2_CONFIG"
	else
		message "Missing tool: $TOOL"
	fi
	return 1
}

is_rootfs()
{
	[ $MOUNT_POINT = "/" ]
}

is_ro()
{
	mountpoint -q $MOUNT_POINT || return 1
	touch $MOUNT_POINT 2>/dev/null || return 0
	return 1
}

normalize_device()
{
	if is_rootfs; then
		DEV_NO="$(mountpoint -d /)"
		ROOT_DEV="$(find /sys/dev/ -name "$DEV_NO" | tail -n 1)"
		echo "/dev/$(basename "$(realpath "$ROOT_DEV")")"
		return
	fi

	case "$DEV" in
		""|tmpfs|proc|devtmpfs|devpts|sysfs|configfs|debugfs|pstore)
			return
			;;
		/dev/*)
			echo $DEV
			return
			;;
	esac

	# Try udev rules symlink
	if find /dev/block/by-name/$DEV 2>/dev/null; then
		return
	fi
	if find /dev/disk/*/$DEV 2>/dev/null; then
		return
	fi

	# Check partition name
	UEVENT=$(grep -l "^PARTNAME=$DEV$" \
		/sys/class/block/*/uevent 2>/dev/null | tail -n 1 || true)
	if [ "$UEVENT" ]; then
		echo "/dev/$(basename $(dirname $UEVENT))"
		return
	fi

	# Check mtdblock name
	for d in $(ls /dev/ | grep mtdblock || true); do
		if grep -q "^$DEV$" /sys/block/$d/device/name; then
			echo "/dev/$d"
			return
		fi
	done

	# Check mtd
	if grep -q "\"$DEV\"$" /proc/mtd 2>/dev/null; then
		echo "/dev/$(grep "\"$DEV\"$" /proc/mtd | cut -d':' -f1)"
		return
	fi

	# Parse from blkid
	if blkid 2>/dev/null | grep -w "$DEV" | tail -n 1 | \
		grep -o "^[^:]*"; then
		return
	fi

	# Give up
	echo $DEV
}

prepare_ubi()
{
	# Only support ubi for mtd device
	if echo $DEV | grep -vq /dev/mtd; then
		message "$DEV is not a mtd device!"
		return 1
	fi

	[ "$PART_NO" ] || { message "No valid part number!" && return 1; }

	MTD_DEV=/dev/mtd${PART_NO}
	UBI_DEV=/dev/ubi${PART_NO}
	UBIVOL_DEV=/dev/ubi${PART_NO}_0
	UBIBLOCK_DEV=/dev/ubiblock${PART_NO}_0

	if [ "$FSGROUP" = ubifs ]; then
		DEV=$UBIVOL_DEV
	else
		DEV=$UBIBLOCK_DEV
	fi

	message "Preparing $DEV from $MTD_DEV"

	message "Remove ubi block device"
	if echo $DEV | grep -q ubiblock; then
		check_tool ubiblock BR2_PACKAGE_MTD_UBIBLOCK || return 1
		ubiblock -r $UBIVOL_DEV >/dev/null 2>/dev/null
	fi

	message "Detach ubi device"
	check_tool ubidetach BR2_PACKAGE_MTD_UBIDETACH || return 1
	ubidetach -p $MTD_DEV >/dev/null 2>/dev/null

	message "Attach ubi device"
	check_tool ubiattach BR2_PACKAGE_MTD_UBIATTACH || return 1
	ubiattach /dev/ubi_ctrl -m $PART_NO -d $PART_NO || return 1

	message "Check for valid volume"
	if [ ! -e $UBIVOL_DEV ]; then
		message "No valid ubi volume"
		return 1
	fi

	message "Create ubi block device"
	if echo $DEV | grep -q ubiblock; then
		check_tool ubiblock BR2_PACKAGE_MTD_UBIBLOCK || return 1
		ubiblock -c $UBIVOL_DEV || return 1
	fi

	return 0
}

format_ubifs()
{
	message "Formatting $MTD_DEV for $DEV"

	message "Remove ubi block device"
	if echo $DEV | grep -q ubiblock; then
		check_tool ubiblock BR2_PACKAGE_MTD_UBIBLOCK || return 1
		ubiblock -r $UBIVOL_DEV >/dev/null 2>/dev/null
	fi

	message "Detach ubi device"
	check_tool ubidetach BR2_PACKAGE_MTD_UBIDETACH || return 1
	ubidetach -p $MTD_DEV >/dev/null 2>/dev/null

	message "Format device"
	check_tool ubiformat BR2_PACKAGE_MTD_UBIFORMAT || return 1
	ubiformat -yq $MTD_DEV || return 1

	message "Attach ubi device"
	ubiattach /dev/ubi_ctrl -m $PART_NO -d $PART_NO || return 1

	message "Create ubi volume"
	check_tool ubimkvol BR2_PACKAGE_MTD_UBIMKVOL || return 1
	ubimkvol $UBI_DEV -N $PART_NAME -m || return 1

	message "Create ubi block device"
	if echo $DEV | grep -q ubiblock; then
		check_tool ubiblock BR2_PACKAGE_MTD_UBIBLOCK || return 1
		ubiblock -c $UBIVOL_DEV || return 1
	fi
}

remount_part()
{
	mountpoint -q $MOUNT_POINT || return

	if is_ro; then
		[ "$1" != ro ] || return
	else
		[ "$1" != rw ] || return
	fi

	mount -o remount,$1 $MOUNT_POINT
}

format_part()
{
	message "Formatting $DEV($FSTYPE)"

	case $FSGROUP in
		ext2)
			# Set max-mount-counts to 0, and disable the time-dependent checking.
			check_tool mke2fs BR2_PACKAGE_E2FSPROGS && \
			mke2fs -F -b 4096 -L $PART_NAME $DEV && \
			tune2fs -c 0 -i 0 $DEV
			;;
		vfat)
			# Use fat32 by default
			check_tool mkfs.vfat BR2_PACKAGE_DOSFSTOOLS_MKFS_FAT && \
			mkfs.vfat -I -F 32 -n $PART_NAME $DEV
			;;
		exfat)
			check_tool mkexfatfs BR2_PACKAGE_EXFAT_UTILS && \
			mkexfatfs -n $PART_NAME $DEV
			;;
		ntfs)
			# Enable compression
			check_tool mkntfs BR2_PACKAGE_NTFS_3G_NTFSPROGS && \
			mkntfs -FCQ -L $PART_NAME $DEV
			;;
		btrfs)
			check_tool mkfs.btrfs BR2_PACKAGE_BTRFS_PROGS && \
			mkfs.btrfs -f -L $PART_NAME $DEV
			;;
		f2fs)
			check_tool mkfs.f2fs BR2_PACKAGE_F2FS_TOOLS && \
			mkfs.f2fs -f -l $PART_NAME $DEV
			;;
		ubifs)
			format_ubifs
			;;
		jffs2)
			check_tool mkfs.jffs2 BR2_PACKAGE_MTD_MKFSJFFS2 && \
			mkfs.jffs2 -o $DEV 0x10000 --pad=0x400000 -s 0x1000 -n
			;;
		erofs|squashfs)
			message "It's pointness to format a read-only partition..."
			false
			;;
		auto)
			message "Unable to format a auto partition..."
			false
			;;
		*)
			message "Unsupported file system $FSTYPE for $DEV"
			false
			;;
	esac
}

format_resize()
{
	BACKUP=$1
	SRC=$(realpath $MOUNT_POINT)

	message "Format-resizing $DEV($FSTYPE)"

	message "Backup original data"
	cp -a "$SRC" "$BACKUP/" || return 1
	umount "$SRC" || return 1

	message "Format and mount rw"
	format_part || return 1
	mount_part || return 1
	remount_part rw

	message "Restore backup data"
	cp -a "$BACKUP/$SRC" $(dirname "$SRC") || return 1
}

resize_ext2()
{
	check_tool resize2fs BR2_PACKAGE_E2FSPROGS_RESIZE2FS || return 1

	resize2fs $DEV
}

resize_vfat()
{
	check_tool fatresize BR2_PACKAGE_FATRESIZE || return 1

	SIZE=$(fatresize -i $DEV | grep "Size:" | grep -o "[0-9]*$")

	# Somehow fatresize only works for 256M+ fat
	[ "$SIZE" -gt $((256 * 1024 * 1024)) ] && return 1

	MAX_SIZE=$(( $(cat $SYS_PATH/size) * 512))
	MIN_SIZE=$(($MAX_SIZE - 16 * 1024 * 1024))
	[ $MIN_SIZE -lt $SIZE ] && return 0 # Large enough!
	while [ $MAX_SIZE -gt $MIN_SIZE ];do
		# Somehow fatresize cannot resize to max size
		MAX_SIZE=$(($MAX_SIZE - 512 * 1024))

		# Try to resize with fatresize, not always work
		fatresize -s $MAX_SIZE $DEV && return
	done
	return 1
}

resize_ntfs()
{
	check_tool ntfsresize BR2_PACKAGE_NTFS_3G_NTFSPROGS || return 1

	echo y | ntfsresize -f $DEV
}

resize_btrfs()
{
	check_tool btrfs BR2_PACKAGE_BTRFS_PROGS || return 1

	btrfs filesystem resize max $DEV
}

resize_f2fs()
{
	check_tool resize.f2fs BR2_PACKAGE_F2FS_TOOLS || return 1

	resize.f2fs $DEV
}

resize_part()
{
	# Already resized
	if [ -f $MOUNT_POINT/.resized ]; then
		message "Already resized $DEV($FSTYPE)"
		return
	fi

	if [ -z "$FSRESIZE" ]; then
		message "No resize for $FSTYPE"
		return
	fi

	message "Resizing $DEV($FSTYPE)"

	# Online resize needs read-write
	remount_part rw
	if eval $FSRESIZE; then
		touch $MOUNT_POINT/.resized
		sync
		return
	fi

	message "Done with rootfs"
	is_rootfs && return

	message "Fallback to format resize"
	TEMP_BACKUP=$(mktemp -d)
	if format_resize $TEMP_BACKUP; then
		touch $MOUNT_POINT/.resized
		sync
	fi
	rm -rf $TEMP_BACKUP
}

convert_mount_opts()
{
	# Accept all opts by default for standard mount tool
	if [ -z "$@" ] && ! mount --help 2>&1 | grep -q BusyBox; then
		echo $OPTS
		return
	fi

	# Filter out unsupported opts
	for opt in ${@:-$BUSYBOX_MOUNT_OPTS}; do
		echo ${OPTS//,/ } | xargs -n 1 | grep -oE "^$opt$"
	done | tr "\n" ","
}

prepare_part()
{
	# Ignore external storages
	case "$MOUNT_POINT" in
		/mnt/* | /media/*) return 1 ;;
	esac

	DEV="$(realpath "$(normalize_device | head -n 1)" 2>/dev/null)"
	[ "$DEV" ] || return 1

	# Only accept block or char device
	[ -b "$DEV" -o -c "$DEV" ] || return 1

	message "Handling $DEV $MOUNT_POINT $FSTYPE $OPTS $PASS"

	SYS_PATH=$(echo /sys/class/*/${DEV##*/})
	if [ -f "$SYS_PATH/name" ]; then
		PART_NAME=$(cat $SYS_PATH/name)
	else
		PART_NAME=$(grep PARTNAME ${SYS_PATH}/uevent | cut -d '=' -f 2)
	fi
	PART_NAME=${PART_NAME:-${DEV##*/}}
	PART_NO=$(echo $DEV | grep -oE "[0-9]*$")

	case $FSTYPE in
		ext[234])
			FSGROUP=ext2
			FSCK="fsck.$FSTYPE -y"
			FSCK_CONFIG=BR2_PACKAGE_E2FSPROGS_FSCK
			FSRESIZE=resize_ext2
			;;
		msdos|fat|vfat)
			FSGROUP=vfat
			FSCK="fsck.vfat -y"
			FSCK_CONFIG=BR2_PACKAGE_DOSFSTOOLS_FSCK_FAT
			FSRESIZE=resize_vfat
			;;
		exfat)
			FSGROUP=exfat
			FSCK="exfatfsck -y"
			FSCK_CONFIG=BR2_PACKAGE_EXFAT_UTILS
			unset FSRESIZE
			;;
		ntfs)
			FSGROUP=ntfs
			FSCK=ntfsfix
			FSCK_CONFIG=BR2_PACKAGE_NTFS_3G_NTFSPROGS
			FSRESIZE=resize_ntfs
			;;
		btrfs)
			FSGROUP=btrfs
			FSCK="btrfs check --repair"
			FSCK_CONFIG=BR2_PACKAGE_BTRFS_PROGS
			FSRESIZE=resize_btrfs
			;;
		f2fs)
			FSGROUP=f2fs
			FSCK="fsck.f2fs -y"
			FSCK_CONFIG=BR2_PACKAGE_F2FS_TOOLS
			FSRESIZE=resize_f2fs
			;;
		ubifs)
			FSGROUP=ubifs
			unset FSCK
			unset FSRESIZE
			;;
		squashfs)
			FSGROUP=squashfs
			unset FSCK
			unset FSRESIZE
			;;
		erofs)
			FSGROUP=erofs
			unset FSCK
			unset FSRESIZE
			;;
		jffs2)
			FSGROUP=jffs2
			unset FSCK
			unset FSRESIZE
			;;
		auto)
			FSGROUP=auto
			message "Running fsck on a random fs is dangerous"
			unset FSCK
			unset FSRESIZE
			;;
		*)
			message "Unsupported file system $FSTYPE for $DEV"
			return 1
	esac

	# Setup mount tool and opts
	case $FSGROUP in
		ntfs)
			if grep -wq ntfs3 /proc/filesystems; then
				MOUNT="mount -t ntfs3"
				OPTS=$(convert_mount_opts)
			else
				MOUNT=ntfs-3g
				check_tool ntfs-3g BR2_PACKAGE_NTFS_3G || return 1
				OPTS=$(convert_mount_opts "$NTFS_3G_MOUNT_OPTS")
			fi
			;;
		auto)
			MOUNT=mount
			OPTS=$(convert_mount_opts)
			;;
		*)
			MOUNT="mount -t $FSTYPE"
			OPTS=$(convert_mount_opts)
			;;
	esac
	MOUNT_OPTS=${OPTS:+" -o ${OPTS%,}"}

	# Detect ro/rw
	MOUNT_RO_RW=$(echo "$OPTS" | tr ',' ' ' | xargs -n 1 | \
		grep -o "^r[ow]$" || true)
	if [ -z "$MOUNT_RO_RW" ]; then
		if is_rootfs && is_ro; then
			MOUNT_RO_RW=ro
		else
			MOUNT_RO_RW=rw
		fi
	fi

	# Prepare for ubi
	if [ "$FSTYPE" = ubifs ] ||
		[ "$(dd if=$DEV bs=1 count=3 status=none)" = UBI ]; then
		if ! prepare_ubi; then
			message "Failed to prepare ubi for $DEV"

			if [ "$AUTO_MKFS" ] && [ "$FSGROUP" = ubifs ]; then
				message "Auto formatting"
				format_ubifs || true
			fi
		fi
	fi
}

check_part()
{
	[ "$SKIP_FSCK" -o "$PASS" -eq 0 ] && return

	if [ -z "$FSCK" ]; then
		message "No fsck for $FSTYPE"
		return
	fi

	message "Checking $DEV($FSTYPE)"

	if is_rootfs; then
		# Remount ro for checking
		remount_part ro
	elif mountpoint -q "$MOUNT_POINT"; then
		# Unmount for checking
		umount "$MOUNT_POINT" || remount_part ro
	fi

	check_tool "$FSCK" $FSCK_CONFIG || return

	$FSCK $DEV
}

mount_part()
{
	mountpoint -q "$MOUNT_POINT" && return

	message "Mounting $DEV($FSTYPE) on $MOUNT_POINT ${MOUNT_OPTS:+with$MOUNT_OPTS}"
	$MOUNT $DEV $MOUNT_POINT $MOUNT_OPTS && return
	[ "$AUTO_MKFS" ] || return

	message "Failed to mount $DEV, try to format it"
	format_part && $MOUNT $DEV $MOUNT_POINT $MOUNT_OPTS
}