#!/bin/sh

source /etc/sysconfig/config;
source /etc/sysconfig/devices;

verbose=0
IMAGE_PATH=.
MTD_PARTS=
start_mtd_dev=0
# this list is order dependent, change with care!
MODULES="mhif chipreg gen_probe cfi_util cfi_probe cfi_cmdset_0002 mtd_blkdevs mtdblock cmdlinepart mtdchar mg3500_flash jffs2"

usage()
{
	echo
	echo "USAGE:  [--path] [--verbose]"
	echo "  --path       - path to images, default is current dir"
	echo "  --verbose    - verbose messages"
	exit
}

get_upgrade_filenames()
{
		UPGRADE_FILENAME_BOOTLOADER=;
		UPGRADE_FILENAME_KERNEL=;
		UPGRADE_FILENAME_INITRD=;
		UPGRADE_FILENAME_DISTPKG=;

		DO_UPGRADE=;
		echo -ne "Upgrade bootloader(y/N): ";
		read DO_UPGRADE
		if [ "X${DO_UPGRADE}" == "Xy" ] || [ "X${DO_UPGRADE}" == "XY" ]
		then
			echo -ne "Enter bootloader filename(nor-bootloader.bin): "
			read UPGRADE_FILENAME_BOOTLOADER
			if [ "X${UPGRADE_FILENAME_BOOTLOADER}" == "X" ] 
			then
				UPGRADE_FILENAME_BOOTLOADER=nor-bootloader.bin
			fi
			if [ ! -e ${IMAGE_PATH}/${UPGRADE_FILENAME_BOOTLOADER} ]
			then
				echo "ERROR: ${IMAGE_PATH}/${UPGRADE_FILENAME_BOOTLOADER} not found"
				exit
			fi
		fi

		DO_UPGRADE=;
		echo -ne "Upgrade kernel(y/N): ";
		read DO_UPGRADE
		if [ "X${DO_UPGRADE}" == "Xy" ] || [ "X${DO_UPGRADE}" == "XY" ]
		then
			echo -ne "Enter kernel filename(nor-kernel.img): "
			read UPGRADE_FILENAME_KERNEL
			if [ "X${UPGRADE_FILENAME_KERNEL}" == "X" ] 
			then
				UPGRADE_FILENAME_KERNEL=nor-kernel.img
			fi
			if [ ! -e ${IMAGE_PATH}/${UPGRADE_FILENAME_KERNEL} ]
			then
				echo "ERROR: ${IMAGE_PATH}/${UPGRADE_FILENAME_KERNEL} not found"
				exit
			fi
		fi

		DO_UPGRADE=;
		echo -ne "Upgrade initrd(y/N): ";
		read DO_UPGRADE
		if [ "X${DO_UPGRADE}" == "Xy" ] || [ "X${DO_UPGRADE}" == "XY" ]
		then
			echo -ne "Enter initrd filename(nor-initrd.img): "
			read UPGRADE_FILENAME_INITRD
			if [ "X${UPGRADE_FILENAME_INITRD}" == "X" ] 
			then
				UPGRADE_FILENAME_INITRD=nor-initrd.img
			fi
			if [ ! -e ${IMAGE_PATH}/${UPGRADE_FILENAME_INITRD} ]
			then
				echo "ERROR: ${IMAGE_PATH}/${UPGRADE_FILENAME_INITRD} not found"
				exit
			fi
		fi

		DO_UPGRADE=;
		echo -ne "Upgrade distpkg(y/N): ";
		read DO_UPGRADE
		if [ "X${DO_UPGRADE}" == "Xy" ] || [ "X${DO_UPGRADE}" == "XY" ]
		then
			while [ "X${UPGRADE_FILENAME_DISTPKG}" == "X" ] 
			do
				echo -ne "Enter distpkg filename: "
				read UPGRADE_FILENAME_DISTPKG
			done
			if [ ! -e ${IMAGE_PATH}/${UPGRADE_FILENAME_DISTPKG} ]
			then
				echo "ERROR: ${IMAGE_PATH}/${UPGRADE_FILENAME_DISTPKG} not found"
				exit
			fi
		fi
}

get_partitions()
{
	REPARTITION="n"
	echo -ne "Re-partition MTD device(y/N/s) ? : "
	read REPARTITION
	if [ "X${REPARTITION}" == "Xs" ]
	then
		echo "CURRENT: ${CFGDEVICES_BOOTDEV_MTDPARTS}"
		echo -ne "Re-partition MTD device(y/N) ? : "
		read REPARTITION
	fi

	if [ "X${REPARTITION}" == "Xy" ] || [ "X${REPARTITION}" == "XY" ]
	then
		REPARTITION="y"
		NEWPARTITIONS=
		echo -ne "Enter partition info: "
		read NEWPARTITIONS
		echo $NEWPARTITIONS |grep -sq '"'
		if [ $? -eq 0 ]
		then
			echo $NEWPARTITIONS |grep -sq '^"'
			if [ $? -eq 0 ]
			then
				NEWPARTITIONS=`echo $NEWPARTITIONS | busybox cut -d'"' -f2`
			else
				NEWPARTITIONS=`echo $NEWPARTITIONS | busybox cut -d'"' -f1`
			fi
		fi
		echo -e "\nWARNING: All data on NOR device will be deleted!!"
		echo -ne "Do you want to continue (y/N) ? "
		read VALIDATE
		if [ "X${VALIDATE}" != "Xy" ] && [ "X${VALIDATE}" != "XY" ]
		then
			echo "Upgrade aborted by user."
			exit 0
		fi
		MTD_PARTS="${NEWPARTITIONS}"

		# XXX for now, just assume that there is only one spare partition
		mount | grep -sq 'on /data type'
		if [ $? -eq 0 ]
		then
			umount /data
		fi

		# unload modules, so we can reload and re-define the partitions
		echo "$MODULES" | awk '{ for (i=NF;i>0;i--) print $i }' | \
		while read modname
		do
			lsmod |grep -sq "$modname[ ]"
			if [ $? -eq 0 ]
			then
				rmmod $modname
			fi
		done

		# hack, we can't unload the mtd core driver which means there are some
		# reference to mtd device still sitting around and we end up with more
		# devices than actually exist, we can detect this before reloading the 
		# modules and thus we know where the "real" devices start!
		if [ -e /proc/mtd ]
		then
			start_mtd_dev=`cat /proc/mtd |tail -1 |cut -f1 -d':' | cut -c4-10`
			start_mtd_dev=$((start_mtd_dev+1));
		fi

		# now re-load, and we should be good to go
		for modname in $MODULES
		do
			if [ "X${modname}" == "Xcmdlinepart" ]
			then
				modprobe ${modname} "mtdparts=physmap-flash:${MTD_PARTS}"
			else
				modprobe ${modname}
			fi
			if [ $? -ne 0 ]
			then
				echo "ERROR: modprobe of ${modname} failed"
				exit
			fi
		done

	else
		MTD_PARTS="${CFGDEVICES_BOOTDEV_MTDPARTS}"
	fi
}

display_filenames()
{
	echo -ne "\nYou requested an upgrade of the following:\n"
	UPGRADE_BOOTLOADER=;
	if [ -n "${UPGRADE_FILENAME_BOOTLOADER}" ]
	then
		echo -ne "Bootloader: ${UPGRADE_FILENAME_BOOTLOADER}\n"
	fi
	if [ -n "${UPGRADE_FILENAME_KERNEL}" ]
	then
		echo -ne "Kernel: ${UPGRADE_FILENAME_KERNEL}\n"
	fi
	if [ -n "${UPGRADE_FILENAME_INITRD}" ]
	then
		echo -ne "Initrd: ${UPGRADE_FILENAME_INITRD}\n"
	fi
	if [ -n "${UPGRADE_FILENAME_DISTPKG}" ]
	then
		echo -ne "Distpkg: ${UPGRADE_FILENAME_DISTPKG}\n"
		UPGRADE_DISTPKG_TYPE="$(echo ${UPGRADE_FILENAME_DISTPKG} | sed -e 's/.*\.//')";
	fi
}

writeFlash()
{
	device=$1
	filename=$2

	flashcp $filename $device
	if [ $? -ne 0 ]
	then
		echo "****************************************************"
		echo "ERROR:  flashcp $filename $device failed..."
		echo "        continuing but system may be un-bootable"
		echo "****************************************************"
	fi
}

info()
{
	echo "==> INFO: $1 <=="
}

verbose()
{
	if [ $verbose -eq 1 ] 
	then
		echo "==> INFO: $1 <=="
	fi
}

while [ $# -ne 0 ]
do
	if [ $1 == "--verbose" ]
	then
		verbose=1;
		shift;
	elif [ $1 == "--path" ]
	then
		IMAGE_PATH=$2
		shift;shift
	else 
		usage
	fi
done

CORRECT=""
until [ "X${CORRECT}" == "Xy" ] || [ "X${CORRECT}" == "XY" ] 
do
	get_upgrade_filenames
	get_partitions
	display_filenames
	if [ "X${REPARTITION}" == "Xy" ]
	then
		echo "Partitions: ${MTD_PARTS}"
	fi
	echo -ne "\nIs this correct(y/N/q) ? ";
	read CORRECT
	if [ "X${VALIDATE}" == "Xq" ]
	then
		exit
	fi
done

echo -e "\nWARNING: All data in selected sections will be deleted!!"
echo -ne "Do you want to continue (y/N) ? ";
read VALIDATE
if [ "X${VALIDATE}" != "Xy" ] && [ "X${VALIDATE}" != "XY" ]
then
	echo "Upgrade aborted by user.";
	exit 0;
fi
echo

extra_start_index=0
OLD_IFS=${IFS}
IFS=","

if [ "X${REPARTITION}" == "Xy" ]
then
	if [ "X${UPGRADE_FILENAME_BOOTLOADER}" == "X" ] || [ "X${UPGRADE_FILENAME_KERNEL}" == "X" ] || [ "X${UPGRADE_FILENAME_INITRD}" == "X" ]
	then
		echo "ERROR: You must update the bootloader, kernel and initrd when repartitioning"
		exit
	fi
fi

i=$start_mtd_dev

for part in ${MTD_PARTS}
do
    name=$(echo ${part} | sed -e 's/[^\(]*[\(]\([^\)]*\)[\)].*/\1/');
    if [ -z "$(echo ${name} | sed -e 's/mboot[0-9]*//')" ]
    then
        type=bootloader
    elif [ -z "$(echo ${name} | sed -e 's/kernel[0-9]*//')" ]
    then
        type=kernel
    elif [ -z "$(echo ${name} | sed -e 's/initrd[0-9]*//')" ]
    then
        type=initrd
    elif [ -z "$(echo ${name} | sed -e 's/rootfs[0-9]*//')" ]
    then
        type=rootfs
    elif [ -z "$(echo ${name} | sed -e 's/data[0-9]*//')" ]
    then
        type=data
	else
		if [ $extra_start_index -eq 0 ]
		then 
			extra_start_index=$i
		fi
        type=extra
    fi

	if [ $type == "bootloader" ] && [ "X${UPGRADE_FILENAME_BOOTLOADER}" != "X" ]
	then
		echo
		info "Upgrading bootloader on /dev/mtd${i}"
		verbose "Erasing /dev/mtd${i}"
		flash_eraseall /dev/mtd${i}
		verbose "Writing /dev/mtd${i}"
		writeFlash /dev/mtd${i} ${IMAGE_PATH}/${UPGRADE_FILENAME_BOOTLOADER}
	elif [ $type == "kernel" ] && [ "X${UPGRADE_FILENAME_KERNEL}" != "X" ]
	then
		echo
		info "Upgrading kernel on /dev/mtd${i}"
		verbose "Erasing /dev/mtd${i}"
		flash_eraseall /dev/mtd${i}
		verbose "Writing /dev/mtd${i}"
		writeFlash  /dev/mtd${i} ${IMAGE_PATH}/${UPGRADE_FILENAME_KERNEL}
	elif [ $type == "initrd" ] && [ "X${UPGRADE_FILENAME_INITRD}" != "X" ]
	then
		echo
		info "Upgrading initrd on /dev/mtd${i}"
		verbose "Erasing /dev/mtd${i}"
		flash_eraseall /dev/mtd${i}
		verbose "Writing /dev/mtd${i}"
		writeFlash /dev/mtd${i} ${IMAGE_PATH}/${UPGRADE_FILENAME_INITRD}
	elif [ $type == "data" ] && [ "X${UPGRADE_FILENAME_DISTPKG}" != "X" ]
	then
		if [ "X${REPARTITION}" == "Xy" ]
		then
			info "Initializing distpkg on /dev/mtd${i}"
			verbose "Erasing /dev/mtd${i}"
			flash_eraseall /dev/mtd${i}
			if [ ! -e /data ]
			then
				mkdir /data
			fi
			mount -t jffs2 /dev/mtdblock${i} /data
		else
			rm -rf /data/distpkg
		fi
		info "Upgrading distpkg on /dev/mtd${i}"
		mkdir /data/distpkg
		if [ -e /data/rootfs ]
		then
			rm -rf /data/rootfs
		fi
		if [ $verbose -eq 1 ]
		then
			tar -zvxf ${IMAGE_PATH}/${UPGRADE_FILENAME_DISTPKG} -C /data
		else
			tar -zxf ${IMAGE_PATH}/${UPGRADE_FILENAME_DISTPKG} -C /data
		fi
		if [ $? -ne 0 ]
		then
			echo "ERROR: tar -zxf ${IMAGE_PATH}/${UPGRADE_FILENAME_DISTPKG} -C /data failed"
			exit
		fi
		sync

		mv /data/rootfs/* /data/distpkg
		rm -rf /data/rootfs
		rm -rf /data/distpkg/etc/ /data/distpkg/lib /data/distpkg/usr/include
	fi
	i=$((i+1));
done
IFS=${OLD_IFS}

echo -e "\nWARNING: You must reboot to enable new software"
echo -ne "reboot now (y/N) ? ";
read VALIDATE
if [ "X${VALIDATE}" == "Xy" ] || [ "X${VALIDATE}" == "XY" ]
then
	reboot
fi
echo


