#!/bin/bash

sdir=`dirname $(readlink -f "${BASH_SOURCE[0]}")`
rootdir=`dirname $sdir`

# codepage is hardcoded in 'loop_rootfs' and scripts, as well as required by pi-gen build
if [ -z ${FAT_CODEPAGE+x} ]; then
    readonly FAT_CODEPAGE="437"
    export FAT_CODEPAGE
fi

# QCOW2 Routines

export CURRENT_IMAGE
export CURRENT_MOUNTPOINT

# 1 boot, 2 data, 3 root
export NBD_DEV
export MAP_BOOT_DEV
export MAP_ROOT_DEV

# set in build.sh
# should be fairly enough for the beginning
#} overwrite here by uncommenting following lines
# BASE_QCOW2_SIZE=15200M

source ${sdir}/lib_nbd_img_handling

###
### Stateful functions ...
###

# find and initialize free block device nodes
init_nbd() {
	if [ -z "${NBD_DEV}" ]; then
        local nbd_dev=$(lib_init_nbd)
        if [ -z "${nbd_dev}" ] ; then
            echo "init_nbd: NBD device not determined. Exit."
            exit 1;
        fi
        NBD_DEV=${nbd_dev}
        MAP_BOOT_DEV=/dev/mapper/${nbd_dev}p1
        MAP_ROOT_DEV=/dev/mapper/${nbd_dev}p2
    fi
}
export -f init_nbd

# connect image with known format to block device
connect_blkdev() {
	init_nbd
    if [ -z "$NBD_DEV" ] ; then
        echo "connect_blkdev: NBD_DEV not determined. Exit."
        exit 1;
    fi
    if ! lib_connect_blkdev3 "${1}" "qcow2" "${NBD_DEV}" ; then
        exit 1;
    fi
	CURRENT_IMAGE="$1"
}
export -f connect_blkdev

# connect raw image to block device
connect_raw_blkdev() {
	init_nbd
    if [ -z "$NBD_DEV" ] ; then
        echo "connect_raw_blkdev: NBD_DEV not determined. Exit."
        exit 1;
    fi
    if ! lib_connect_blkdev3 "${1}" "raw" "${NBD_DEV}" ; then
        exit 1;
    fi
	CURRENT_IMAGE="$1"
}
export -f connect_raw_blkdev

# disconnect image from block device
disconnect_blkdev() {
    if [ -z "$NBD_DEV" ] ; then
        echo "disconnect_blkdev: NBD_DEV not set. Exit."
        exit 1;
    fi
    lib_disconnect_blkdev "${NBD_DEV}"
	NBD_DEV=
	MAP_BOOT_DEV=
	MAP_ROOT_DEV=
	CURRENT_IMAGE=
}
export -f disconnect_blkdev

# mount qcow2 image: mount_image <image file> <mountpoint>
mount_qimage() {
    if [ -z "$1" ] ; then
        echo "mount_qimage: image-file not given. Exit."
        exit 1;
    fi
    if [ -z "$2" ] ; then
        echo "mount_qimage: mountpoint not given. Exit."
        exit 1;
    fi

	connect_blkdev "$1"

    if [ -z "$NBD_DEV" ] ; then
        echo "mount_qimage: NBD_DEV not determined. Exit."
        exit 1;
    fi
    if [ -z "$MAP_ROOT_DEV" -o -z "$MAP_BOOT_DEV" ] ; then
        echo "mount_qimage: map devices not determined. Exit."
        exit 1;
    fi
	mount -v -t ext4 $MAP_ROOT_DEV "$2"
	mkdir -p "$2/boot"
	mount -v -t ${BOOT_FSTYPE} $MAP_BOOT_DEV "$2/boot"
	CURRENT_MOUNTPOINT="$2"
}
export -f mount_qimage

# mount raw image: mount_image <image file> <mountpoint>
mount_rawimage() {
    if [ -z "$1" ] ; then
        echo "mount_rawimage: image-file not given. Exit."
        exit 1;
    fi
    if [ -z "$2" ] ; then
        echo "mount_rawimage: mountpoint not given. Exit."
        exit 1;
    fi

	connect_raw_blkdev "$1"

    if [ -z "$NBD_DEV" ] ; then
        echo "mount_rawimage: NBD_DEV not determined. Exit."
        exit 1;
    fi
    if [ -z "$MAP_ROOT_DEV" -o -z "$MAP_BOOT_DEV" ] ; then
        echo "mount_rawimage: map devices not determined. Exit."
        exit 1;
    fi
	mount -v -t ext4 $MAP_ROOT_DEV "$2"
	mkdir -p "$2/boot"
	mount -v -t ${BOOT_FSTYPE} $MAP_BOOT_DEV "$2/boot"
	CURRENT_MOUNTPOINT="$2"
}
export -f mount_rawimage

# umount qcow2 or raw image: umount_image <current mountpoint>
umount_image() {
	sync
    if [ -z "$1" ] ; then
        echo "umount_image: mountpoint not given. Exit."
        exit 1;
    fi
	#umount "$1/boot"
	while mount | grep -q "$1"; do
		local LOCS
		LOCS=$(mount | grep "$1" | cut -f 3 -d ' ' | sort -r)
		for loc in $LOCS; do
			echo "$loc"
			while mountpoint -q "$loc" && ! umount "$loc"; do
				sleep 0.1
			done
		done
	done
	CURRENT_MOUNTPOINT=
	disconnect_blkdev
}
export -f umount_image

# create base image / backing image / mount image
load_qimage() {
	if [ -z "${CURRENT_MOUNTPOINT}" ]; then
		if [ ! -d "${ROOTFS_DIR}" ]; then
			mkdir -p "${ROOTFS_DIR}";
		fi

		if [ "${CLEAN}" = "1" ] && [ -f "${WORK_DIR}/image-${STAGE}.qcow2" ]; then
			rm -f "${WORK_DIR}/image-${STAGE}.qcow2";
		fi

		if [ ! -f "${WORK_DIR}/image-${STAGE}.qcow2" ]; then
			pushd ${WORK_DIR} > /dev/null
			init_nbd
			if [ -z "${PREV_STAGE}" ]; then
				echo "Creating base image: image-${STAGE}.qcow2"
				#  -o preallocation=falloc
				qemu-img create -f qcow2 image-${STAGE}.qcow2 $BASE_QCOW2_SIZE
				sync
				qemu-nbd --discard=unmap -c /dev/${NBD_DEV} image-${STAGE}.qcow2
				sync

                if [ "${ROOTFS_RO}" = "1" ] ; then
                    # 1 boot (7100 MiB), rootfs (remainder of BASE_QCOW2_SIZE)
                    sfdisk /dev/${NBD_DEV} << EOF
4MiB,7100MiB,c,*
7104MiB,,83;
EOF
                else
                    # 1 boot (250 MiB), rootfs (remainder of BASE_QCOW2_SIZE)
                    sfdisk /dev/${NBD_DEV} << EOF
4MiB,250MiB,c,*
254MiB,,83;
EOF
                fi
				sync

				kpartx -a /dev/${NBD_DEV}
                case ${BOOT_FSTYPE} in
                    xfs)
                        mkfs.xfs -L BOOT $MAP_BOOT_DEV
                        ;;
                    ext4)
                        mkfs.ext4 -L BOOT -O "^huge_file,^metadata_csum,^64bit" $MAP_BOOT_DEV
                        ;;
                    ext2)
                        mkfs.ext2 -L BOOT -O "^has_journal,^huge_file,^metadata_csum,^64bit" $MAP_BOOT_DEV
                        ;;
                    vfat)
                        mkfs.fat -n BOOT -F 32 --codepage=${FAT_CODEPAGE} -v $MAP_BOOT_DEV
                        ;;
                    *)
                        echo "BOOT_FSTYPE not one of[xfs,ext4,ext2,vfat] but '${BOOT_FSTYPE}'. Exit."
                        exit 1;
                        ;;
                esac
                    
                if [ "${ROOTFS_RO}" = "1" ] ; then
                    # Set reserved-blocks-percentage for root to zero, no journaling
                    mkfs.ext4 -L ROOTFS -m 0 -O "^has_journal,^huge_file,^metadata_csum,^64bit" $MAP_ROOT_DEV
                else
                    mkfs.ext4 -L ROOTFS -O "^huge_file,^metadata_csum,^64bit" $MAP_ROOT_DEV
                fi
				sync
			else
				if [ ! -f "${WORK_DIR}/image-${PREV_STAGE}.qcow2" ]; then
					exit 1;
				fi
				echo "Creating backing image: image-${STAGE}.qcow2 <- ${WORK_DIR}/image-${PREV_STAGE}.qcow2"
				qemu-img create -f qcow2 \
					-o backing_file=${WORK_DIR}/image-${PREV_STAGE}.qcow2 \
					${WORK_DIR}/image-${STAGE}.qcow2
				sync
				qemu-nbd --discard=unmap -c /dev/${NBD_DEV} image-${STAGE}.qcow2
				sync
				kpartx -a /dev/${NBD_DEV}
			fi

			mount -v -t ext4 $MAP_ROOT_DEV "${ROOTFS_DIR}"
			mkdir -p "${ROOTFS_DIR}/boot"
			mount -v -t ${BOOT_FSTYPE} $MAP_BOOT_DEV "${ROOTFS_DIR}/boot"
			CURRENT_IMAGE=${WORK_DIR}/image-${STAGE}.qcow2
			CURRENT_MOUNTPOINT=${ROOTFS_DIR}
			popd > /dev/null
		else
			mount_qimage "${WORK_DIR}/image-${STAGE}.qcow2" "${ROOTFS_DIR}"
		fi
		echo "Current image in use: ${CURRENT_IMAGE} (MP: ${CURRENT_MOUNTPOINT})"
	fi
}
export -f load_qimage

# umount current image and refresh mount point env var
unload_qimage() {
	if [ ! -z "${CURRENT_MOUNTPOINT}" ]; then
		fstrim -v "${CURRENT_MOUNTPOINT}" || true
		umount_image "${CURRENT_MOUNTPOINT}"
	fi
}
export -f unload_qimage

function e2fs_block_size() {
    local res
	res=$(dumpe2fs -h $1 | grep 'Block size' | awk -F': ' ' { print $2 }')
	res=${res// /}
    echo $res
}
export -f e2fs_block_size

# based on: https://github.com/SirLagz/RaspberryPi-ImgAutoSizer
# helper function for make_bootable_image, do not call directly
function resize_qcow2() {
	if [ -z "$CALL_FROM_MBI" ]; then
		echo "resize_qcow2: cannot be called directly, use make_bootable_image instead"
		return 1
	fi

    # 1 boot, 2 rootfs
    #
    # rootfs ext4 ^journal => ext2, second partition
	# ROOT_MARGIN=$((800*1024*1024))
	ROOT_MARGIN=$((1*1024*1024))
	PARTED_OUT=`parted -s -m "/dev/${NBD_DEV}" unit B print`
	#PART_NO=`echo "$PARTED_OUT" | grep ext[24] | awk -F: ' { print $1 } '`
	#PART_START=`echo "$PARTED_OUT" | grep ext[24] | awk -F: ' { print substr($2,1,length($2)-1) } '`
    PART_NO=2
	PART_START=`echo "$PARTED_OUT" | grep "^${PART_NO}:" | awk -F: ' { print substr($2,1,length($2)-1) } '`

	e2fsck -y -f $MAP_ROOT_DEV || true

	DATA_SIZE=`resize2fs -P $MAP_ROOT_DEV | awk -F': ' ' { print $2 } '`
	BLOCK_SIZE=$(e2fs_block_size $MAP_ROOT_DEV)

	let DATA_SIZE=$DATA_SIZE+$ROOT_MARGIN/$BLOCK_SIZE
	resize2fs -p $MAP_ROOT_DEV $DATA_SIZE
	sleep 1

	let PART_NEW_SIZE=$DATA_SIZE*$BLOCK_SIZE
	let PART_NEW_END=$PART_START+$PART_NEW_SIZE
    ACT1=`parted -s "/dev/${NBD_DEV}" rm ${PART_NO}`
    ACT2=`parted -s "/dev/${NBD_DEV}" unit B mkpart primary $PART_START $PART_NEW_END`
    NEW_IMG_SIZE=`parted -s -m "/dev/${NBD_DEV}" unit B print free | tail -1 | awk -F: ' { print substr($2,1,length($2)-1) } '`
    ROOT_PART_START=$PART_START
    ROOT_PART_SIZE_BYTES=$PART_NEW_SIZE
    ROOT_PART_BLOCK_SIZE=$BLOCK_SIZE
    ROOT_PART_BLOCK_COUNT=$DATA_SIZE
}
export -f resize_qcow2

function write_partitions_qcow2() {
	if [ -z "$CALL_FROM_MBI" ]; then
		echo "write_partitions_qcow2: cannot be called directly, use make_bootable_image instead"
		return 1
	fi
    local map_dev
    local block_size
    local block_count
    local img_file
    local inf_file
    local byte_size
    map_dev=$1
    block_size=$2
    block_count=$3
    img_file=$4
    inf_file=$5

    # typical block sizes:
    #   1 boot (512), 2 rootfs (4096)
    #
    if [ "$block_count" = "0" ] ; then
        # copy whole partition
        echo "Copy partition image $map_dev -> image $img_file with blocksize $block_size" | tee -a $inf_file
        dd if=$map_dev of=$img_file bs=$block_size 2>&1 | tee -a $inf_file
    else
        # copy block_count blocks only (smaller filesystem on bigger partition for rootfs)
        let byte_size=$block_count*$block_size
        echo "Size $byte_size B = $block_size B/block * $block_count blocks" >> $inf_file

        echo "Copy partition image $map_dev -> image $img_file with $block_count x blocksize $block_size" | tee -a $inf_file
        dd if=$map_dev of=$img_file bs=$block_size count=$block_count 2>&1 | tee -a $inf_file
    fi
}
export -f write_partitions_qcow2

# create raw img from qcow2: make_bootable_image <in.qcow2> <out.img>
function make_bootable_image() {
	EXPORT_QCOW2="$1"
	EXPORT_IMAGE="$2"
	EXPORT_IMAGE_ROOT="$3"
	EXPORT_INFO_ROOT="$4"

	echo "Connect block device to source qcow2"
	connect_blkdev "${EXPORT_QCOW2}"
    
	CALL_FROM_MBI=1
	echo "Resize fs"
	resize_qcow2
	sync
	CALL_FROM_MBI=
    
	echo "Disconnect block device"
	disconnect_blkdev
    
	if [ -z "$NEW_IMG_SIZE" ]; then
		echo "NEW_IMG_SIZE could not be calculated, cannot process image. Exit."
		exit 1
	fi
	echo "Shrinking qcow2 image"
	qemu-img resize --shrink "${EXPORT_QCOW2}" $NEW_IMG_SIZE
	sync

	echo "Convert qcow2 to raw image"
	qemu-img convert -f qcow2 -O raw "${EXPORT_QCOW2}" "${EXPORT_IMAGE}"
	sync

	echo "Get PARTUUIDs from image"
	IMGID="$(blkid -o value -s PTUUID "${EXPORT_IMAGE}")"

	echo "Mount image ${EXPORT_IMAGE}"
	MOUNTROOT=${WORK_DIR}/tmpimage
	mkdir -p $MOUNTROOT
    mount_rawimage "${EXPORT_IMAGE}" $MOUNTROOT

	if [ ! -d "${MOUNTROOT}/root" ]; then
		echo "Image damaged or not mounted. Exit."
		exit 1
	fi

    if [ "${ROOTFS_RO}" = "1" ] ; then
        echo "Setup hard devicenames"
        BOOT_DEV_NAME="\/dev\/mmcblk0p1"
        echo "Boot: $BOOT_DEV_NAME"
        ROOT_DEV_NAME="\/dev\/mmcblk0p2"
        echo "Root1: $ROOT_DEV_NAME"
        
        echo "Set hard device names in fstab (rootfs ro -> /boot uncommented)"
        sed -i "s/BOOTDEV/${BOOT_DEV_NAME}/" "${MOUNTROOT}/etc/fstab"
        sed -i "s/ROOTDEV/${ROOT_DEV_NAME}/" "${MOUNTROOT}/etc/fstab"
        if [ "${TARGET_RASPI}" = "1" ]; then
            echo "Set rootfs ro hard device name in cmdline, required"
            sed -i "s/ROOTDEV/${ROOT_DEV_NAME}/" "${MOUNTROOT}/boot/sys_${TARGET_ARCH}_000/cmdline.txt"
        fi
    else
        echo "Setup PARTUUIDs"
        BOOT_PARTUUID="${IMGID}-01"
        echo "Boot: $BOOT_PARTUUID"
        ROOT_PARTUUID="${IMGID}-02"
        echo "Root1: $ROOT_PARTUUID"

        if [ ! -z "$BOOT_PARTUUID" ] && [ ! -z "$ROOT_PARTUUID" ]; then
            echo "Set UUIDs to make it bootable, required."
            sed -i "s/BOOTDEV/PARTUUID=${BOOT_PARTUUID}/" "${MOUNTROOT}/etc/fstab"
            sed -i "s/ROOTDEV/PARTUUID=${ROOT_PARTUUID}/" "${MOUNTROOT}/etc/fstab"
            if [ "${TARGET_RASPI}" = "1" ]; then
                sed -i "s/ROOTDEV/PARTUUID=${ROOT_PARTUUID}/" "${MOUNTROOT}/boot/sys_${TARGET_ARCH}_000/cmdline.txt"
            fi
        fi
    fi

	echo "Umount image ${EXPORT_IMAGE}"
	sync
    umount_image $MOUNTROOT

	echo "Extracting partitions - Writing partitions from ${EXPORT_IMAGE}"
	connect_raw_blkdev "${EXPORT_IMAGE}"
    
	CALL_FROM_MBI=1
    # 1 boot (512), 2 rootfs (4096)
	echo "Extract partitions to raw image"
	write_partitions_qcow2 $MAP_ROOT_DEV  $ROOT_PART_BLOCK_SIZE $ROOT_PART_BLOCK_COUNT $EXPORT_IMAGE_ROOT $EXPORT_INFO_ROOT
	sync
	CALL_FROM_MBI=
    
	echo "Disconnect block device"
	disconnect_blkdev
    
	echo "Remove qcow2 export image"
	rm -f "${EXPORT_QCOW2}"
}
export -f make_bootable_image