#!/bin/bash
#
# Reference: <<i.MX Linux® User's Guide.pdf>>
#
# SD Card partition layout:   1 Sector= 512B
#
# +-------------------------------+-----------------------+---------------------------------+
# |         Start  Address        |         Size          |           Usage                 |
# +-------------------------------+-----------------------+---------------------------------+
# |             0x0               |     2 sectors(1K)     |  Reserved for partition table   |
# +-------------------------------+-----------------------+---------------------------------+
# |      2 sector(1K, 0x400)      |  20414 sectors(9M+)   |     i.MX6ULL  u-boot image      |
# +-------------------------------+-----------------------+---------------------------------+
# |  20480 sector(10M, 0xa00000)  | 204800 sectors(100MB) |      FAT32 Boot partition       |
# +-------------------------------+-----------------------+---------------------------------+
# | 225280 sector(110M,0x6e00000) |    Remaining Space    |     EXT4 filesystem for rootfs  |
# +-------------------------------+-----------------------+---------------------------------+
#

PRJ_PATH=`pwd`
PRJ_NAME=`basename ${PRJ_PATH}`

IMAGE_PATH=${PRJ_PATH}/boot/

BOARD=igkboard

# rootfs should be buildroot or stretch/buster/bullseye for debian system
ROOTFS=buildroot

# linux kernel tarball path and branch
TAR_PATH=${PRJ_PATH}/../tarball

BOOT=emmc

IMAGE_NAME=linuxsys_${BOARD}_${ROOTFS}.img
if [ $ROOTFS == buildroot ] ; then
    IMAGE_SIZE=512
else
    IMAGE_SIZE=2048
fi

#i.MX6/i.MX7 read uboot from mmc start on sector #2
UBOOT_OFSET=2
MMC_SECSIZE=512

# U-boot space Size 10MB
UBOOT_SIZE=10

# vfat boot partition 100MB
BOOT_SIZE=100


IMG_UBOOT=${PRJ_PATH}/u-boot-${BOARD}-emmc.imx
IMG_KERNEL=${IMAGE_PATH}/zImage
IMG_DTB=${IMAGE_PATH}/${BOARD}-emmc.dtb
ROOTFS_TAR=${PRJ_PATH}/../rootfs/rootfs_${ROOTFS}.tar.bz2

MNT_POINT=./mnt
LOOP_DEV=`losetup  -f | cut -d/ -f3`

set -u
set -e

function exit_handler()
{
    echo "Shell script exit now, do some clean work"
    echo ""

    set +e

    mountpoint $MNT_POINT > /dev/null 2>&1
    if [ $? == 0 ] ; then
        echo "umount ${MNT_POINT}"
        umount ${MNT_POINT}
    fi

    rm -rf ${MNT_POINT}

    if [ -L /dev/mapper/${LOOP_DEV}p1 ] ; then
        echo "kpartx -dv /dev/${LOOP_DEV}"
        kpartx -dv /dev/${LOOP_DEV}
    fi

    losetup -a | grep "${LOOP_DEV}" > /dev/null 2>&1
    if [ $? == 0 ]  ; then
        echo "losetup -d /dev/${LOOP_DEV}"
        losetup -d /dev/${LOOP_DEV}
    fi
}

STAGE=0
function msg_banner()
{
    STAGE=`expr $STAGE + 1`

    echo ""
    echo "+---------------------------------------------+"
    printf " Stage $STAGE: $1\n"
    echo "+---------------------------------------------+"
    echo ""
}


function generate_image()
{
    # FAT32 boot partition start/end address in MB
    BOOT_START=${UBOOT_SIZE}
    BOOT_END=`expr ${BOOT_START} + ${BOOT_SIZE}`

    msg_banner " <${STAGE}> Generate system image "

    dd if=/dev/zero of=${IMAGE_NAME} bs=1024k count=${IMAGE_SIZE}  && sync
    chmod a+x ${IMAGE_NAME}

    msg_banner " <${STAGE}> Partition system image "

    parted ${IMAGE_NAME} mklabel msdos
    parted ${IMAGE_NAME} mkpart primary fat32 ${BOOT_START}M ${BOOT_END}M
    parted ${IMAGE_NAME} mkpart primary ext4 ${BOOT_END}M 100%

    sync
}


function format_partition()
{
    msg_banner " <${STAGE}> losetup image on $LOOP_DEV"

    losetup /dev/${LOOP_DEV}  ${IMAGE_NAME}
    kpartx -av /dev/${LOOP_DEV}

    msg_banner " <${STAGE}> Format system image partition"
    mkfs.vfat /dev/mapper/${LOOP_DEV}p1
    mkfs.ext4 /dev/mapper/${LOOP_DEV}p2
    sync
}

function install_sysimg()
{
    msg_banner " <${STAGE}> Install u-boot image"
    sudo dd if=${IMG_UBOOT} of=${IMAGE_NAME} bs=${MMC_SECSIZE} seek=${UBOOT_OFSET} conv=notrunc,sync

    msg_banner " <${STAGE}> Install linux kernel image"

    mkdir -p ${MNT_POINT}
    mount -t vfat /dev/mapper/${LOOP_DEV}p1 ${MNT_POINT}

    set -x
    cp ${IMG_KERNEL} ${MNT_POINT}
    cp ${IMG_DTB}    ${MNT_POINT}
    sync
    set +x

    umount ${MNT_POINT}
}

function install_rootfs()
{
    msg_banner " <${STAGE}> Install root filesystem "

    if [ ! -f ${ROOTFS_TAR} ] ; then
        echo "ERROR: rootfs tarball file '${ROOTFS_TAR}' not exist..."
        exit ;
    fi

    mount -t ext4 /dev/mapper/${LOOP_DEV}p2 ${MNT_POINT}
    tar -xjf ${ROOTFS_TAR} -C ${MNT_POINT} && sync
    chown -R root.root ${MNT_POINT} && sync

    umount ${MNT_POINT}
}

function do_distclean()
{
    printf "\n\n -- do distclean in `basename ${PRJ_PATH}` --\n\n"

    rm -f  *${BOARD}*
    rm -rf boot/*Image*
    rm -rf boot/overlays/
    rm -rf boot/*${BOARD}*
    exit 0;
}

function do_image()
{
    mkdir -p ${MNT_POINT}

    generate_image

    format_partition

    install_sysimg

    install_rootfs

    msg_banner " bzip2 compress system image "
    rm -f ${IMAGE_NAME}*.bz2
    bzip2 ${IMAGE_NAME}

    rm -rf ${MNT_POINT}

    printf "\n\n -- generate system image done --\n\n"
}

function do_root()
{
    echo ""
    if [[ $1 == "yes" ]] && [ `id -u` != 0 ] ; then
        echo "ERROR: This action must run as root!"
        echo ""
        exit;
    elif [[ $1 != "yes" ]] && [ `id -u` == 0 ] ; then
        echo "ERROR: This action cannot run as root!"
        echo ""
        exit;
    fi
}

function do_usage()
{
    echo ""
    echo "Usage:"
    echo "   $0 [-b] [-c] [-h]"
    echo "       -b: download and build $PRJ_NAME"
    echo "       -c: clean all the source code"
    echo "       -h: show this help message"
    echo ""
    echo " WARNNING: This shell script must run as sudo"
    echo ""
    exit;
}

trap 'exit_handler' EXIT

while getopts "bch" OPTNAME
do
    case "${OPTNAME}" in
        "b")
            break;
            ;;

        "c")
            do_root "yes"
            do_distclean
            ;;

        "*")
            do_usage
            ;;
    esac
done

do_root "yes"
do_image
exit;