#!/bin/bash
VERSION=1.0

THIS_COMPUTER_IP="192.168.0.230"
UBOOT_IMG_NAME="apf6-u-boot.img"
UBOOT_SPL_NAME="apf6-u-boot.spl"
TFTP_SERVER_PATH="/tftpboot"
BLOCK_SIZE=0x200
RED='\033[0;31m'
NC='\033[0m' # No Color

images=""
#images="apf6-boot.ext4;1 apf6-rootfs.ext4;2"
#images="boot.ext4;1 rootfs.ext2;2"
disk=""
tty=""
ubootimg=""
ums=0

set -euo pipefail

function waitForText
{
        while read input
        do
                if [[ "$input" == *"$2"* ]]; then
			echo "$input"
                        break;
                fi
        done < "$1"
}

function enterUboot
{
        set +e
        fuser "$1" > /dev/null 2>&1
        if [ $? -eq 0 ]; then
                >&2 echo "$1 is in use, exiting"
                exit 1
        fi
        set -e
        stty -F "$1" 115200 raw -echo -echoe -echok #-echoctl -echoke
        while read input
        do
                if [[ "$input" == *"BIOS>" ]]; then
                        break;
                fi
                echo "" > "$1"
        done < "$1"
}

function waitboot
{
	waitForText "$1" "U-Boot SPL" > /dev/null
}

function flash_uboot
{
        if [ ! -e "$2" ]; then
                >&2 echo "Cannot access $2 img, exiting"
                exit 1
        fi

        spl_img=${2%.*}.spl
        if [ ! -e "$spl_img" ]; then
                >&2 echo "Cannot access $spl_img img, exiting"
                exit 1
        fi

        flashDiskNet "$1" "$spl_img;4 $2;3"
	echo "Done"
        echo "Please reboot the board"
	waitboot "$1"
	enterUboot "$1"
        echo "Setting u-boot env"
        echo "run flash_reset_env" > "$1"
	sleep 1
        echo "setenv extrabootargs \"quiet\"; setenv bootdelay 0; setenv bootcmd \"run loadbootscript; run bootscript;\"; setenv bootscript \"echo running bootscript; source \${scriptaddr}\"; saveenv" > "$1"
	sleep 1
        echo "Done"
}

function startUMS
{
        echo " " > "$1"
        echo " " > "$1"
        echo "ums 0 mmc 0" > "$1"
}

function stopUMS
{
        expect -c "send \003;" > "$1"
}

function detectUMS
{
        sleep 3
        umsDisk=$(lsblk -o NAME,MODEL | grep "UMS disk" | cut -d' ' -f1)
        nbUms=$(lsblk -o NAME,MODEL | grep "UMS disk" | cut -d' ' -f1 | wc -l)
        #nbUms=$(echo -n "$umsDisk"| wc -l)
        if [ "$nbUms" != "1" ]; then
                >&2 echo "Cannot auto detect disk, exiting"
                exit 1
        fi

        echo "$umsDisk"
}

function flashDiskNet
{
        local tty="$1"
        local images="$2"
        if [ "$tty" == "" ]; then
                echo "Flash disk by eth need tty option. Giving Up"
                exit 1
        fi
        echo "images to flash $images"
        echo "setenv ipaddr 192.168.0.100; setenv serverip ${THIS_COMPUTER_IP}" > "$tty"
        sleep 0.5
        for desc in $images; do
                image=$(echo "$desc" | cut -d';' -f1)
                part=$(echo "$desc" | cut -d';' -f2)
                crc_file=$(crc32 "$image" | cut -f1)
                echo "Flashing $image with CRC ${crc_file}"
                case "$part" in
                        1)
                                cp "$image" $TFTP_SERVER_PATH/apf6-boot.ext4
                                echo "run update_boot" > "$tty"
                                waitForText "$tty" "Flashing of boot image succeed" > /dev/null
                                echo "mmc read \${loadaddr} 0x800 \${nbblocks}; crc32 \${loadaddr} \${filesize}" > "$tty"
                                ;;
                        2)
                                cp "$image" $TFTP_SERVER_PATH/apf6-rootfs.ext4
                                echo "run update_rootfs" > "$tty"
                                waitForText "$tty" "Flashing of rootfs image succeed" > /dev/null
                                echo "mmc read \${loadaddr} 0x18800 \${nbblocks}; crc32 \${loadaddr} \${filesize}" > "$tty"
                                ;;
                        3)
                                cp "$image" $TFTP_SERVER_PATH/$UBOOT_IMG_NAME
                                echo "run download_uboot_img flash_uboot_img" > "$tty"
                                waitForText "$tty" "Flashing of" > /dev/null
                                echo "mmc read \${loadaddr} 0x8a \${nbblocks}; crc32 \${loadaddr} \${filesize}" > "$tty"
                                ;;
                        4)
                                cp "$image" $TFTP_SERVER_PATH/$UBOOT_SPL_NAME
                                echo "run download_uboot_spl flash_uboot_spl" > "$tty"
                                waitForText "$tty" "Flashing of" > /dev/null
                                echo "mmc read \${loadaddr} 0x2 \${nbblocks}; crc32 \${loadaddr} \${filesize}" > "$tty"
                                ;;
                        *)
                                echo "Unsupported partition $part"
                                exit 1
                                ;;
                esac
                crc_computed=$(waitForText "$tty" "CRC32 for" | awk -F "==>" '{print $2}' | tr -d '[:space:]')
                echo "Computed checksum ${crc_computed}"
                if [ "$crc_computed" != "$crc_file" ]; then
                        echo -e "${RED}CHECKSUM ERROR${NC} for $image (got \"$crc_computed\" but should be \"$crc_file\")"
                fi

                echo "done"
        done
}

function flashDisk
{
        local tty="$1"
        local images="$2"
	echo "images to flash $images"
	if [ "$tty" != "" ]; then
		echo "Start UMS"
		startUMS "$tty"
	fi

	if [ "$disk" == "" ]; then
		disk=$(detectUMS)
	fi

	if [ "$disk" == "" ]; then
		>&2 echo "Cannot found disk, exiting"
		exit 1
	fi

	if [ ! -e "/dev/$disk" ]; then
		>&2 echo "disk /dev/$disk not found. Giving up"
		exit 1
	fi

	for desc in $images; do
		image=$(echo "$desc" | cut -d';' -f1)
		part=$(echo "$desc" | cut -d';' -f2)
		diskPath="/dev/${disk}${part}"
		echo "flashing $image on $diskPath"
		if [ ! -e "$image" ]; then
			>&2 echo "$image cannot be found"
			exit 1
		fi
		if grep -qs $diskPath /proc/mounts; then
			echo "$diskPath mounted ! unmouting it"
			echo "Is disk automounted ? try to run \"gsettings set org.gnome.desktop.media-handling automount false\" to disable it"
			umount $diskPath

		fi
		sudo dd if="$image" of="$diskPath" bs=4M status=progress oflag=direct
		# flush file system buffers. Force changed blocks to disk, update the super block.
		sync
	done

        if [ "$tty" != "" ]; then
                stopUMS "$tty"
                sleep 0.5
                for desc in $images; do
                        image=$(echo "$desc" | cut -d';' -f1)
                        part=$(echo "$desc" | cut -d';' -f2)
                        crc_file=$(crc32 "$image" | cut -f1)
                        size=$(du -b "$image" | cut -f1)
                        size_hex=$(printf "%x" "$size")
                        part_addr=0
                        echo "Computing flashed $image CRC ${crc_file} of size 0x${size_hex}"
                        nbblocks=$(( 0x$size_hex / $BLOCK_SIZE + 1))
                        nbblocks=$(printf "%x" "$nbblocks")
                        if [ "$part" == "1" ]; then
                                part_addr=0x800
                        else
                                part_addr=0x18800
                        fi
                        echo "mmc read \${loadaddr} ${part_addr} 0x${nbblocks}; crc32 \${loadaddr} 0x${size_hex}" > "$tty"
                        crc_computed=$(waitForText "$tty" "CRC32 for" | awk -F "==>" '{print $2}' | tr -d '[:space:]')
                        echo "Computed checksum ${crc_computed}"
                        if [ "$crc_computed" != "$crc_file" ]; then
                                echo -e "${RED}CHECKSUM ERROR${NC} for $image (got $crc_computed but should be $crc_file)"
				exit 1
                        fi
                done
        else
                echo "no tty given -> cannot computed CRC"
        fi

	if [ "$tty" != "" ]; then
		echo "Restart board for first boot"
                echo "reset" > "$tty"
	fi
	# wait for first boot to complete
	sleep 30
}

echo "Version: $VERSION"
while getopts "afb:s:d:t:u:" opt; do
        case $opt in
                a)      echo "Flashing all images"
                        ;;
                b)      echo "Flashing boot image ${OPTARG}"
                        images="${images} ${OPTARG};1"
                        ;;
                s)      echo "Flashing system image ${OPTARG}"
                        images="${images} ${OPTARG};2"
                        ;;
                d)      echo "Using disk $OPTARG"
                        disk=$OPTARG
                        ;;
                t)      echo "Using tty $OPTARG"
                        tty=$OPTARG
                        ;;
                u)      echo "Flashing uboot $OPTARG"
                        ubootimg=$OPTARG
                        ;;
		f)	echo "use UMS for flashing"
			ums=1
			;;
                \?)
                        echo "Invalid option: -$opt" >&2
                        exit 1
                        ;;
        esac
done


if [ "$tty" == "" ] && [ "$ubootimg" != "" ]; then
        >&2 echo "TTY should be provided when flashing uboot"
        exit 1
fi

if [ "$tty" != "" ]; then
        echo "Please reboot the board"
	enterUboot "$tty"
fi

if [ "$ubootimg" != "" ]; then
	flash_uboot "$tty" "$ubootimg"
fi

if [ "$images" != "" ]; then
	if [ $ums == 1 ]; then
		flashDisk "$tty" "$images"
	else
		flashDiskNet "$tty" "$images"
	fi
fi

echo "SOM programmation done"
notify-send --urgency=low "Flashing PRIMA2 finished"