#!/bin/sh

source /etc/sysconfig/config;
source /etc/sysconfig/devices;
source /etc/sysconfig/helpers/functions

usage() 
{

	echo
	echo "USAGE:  --parts <quoted partition def> [--nand | --nor] [--path] [--verbose][--skip_extra][--help] "
	echo "	partition definition must be supplied enclosed in '\"'"
	echo "	--nand is the default"
	echo "	--nor is supported if the size is  16MB "
	echo "	--path - path to images, default is current dir"
	echo "	--skip_extra - do not initialize extra, non-core partitions"
	echo "	--init_extra - define multiple extra partitions, see verbose help"
	echo "	give --help --verbose for vebose help"
	echo

	if [ $verbose -eq 1 ]
	then
		echo "Verbose help:"
		echo "  --skip_extra: core partitions are the mboot, initrd, kernel and rootfs"
		echo "partitions, using this switch will skip all initialization(including erasing)"
		echo "on any partitions beyond these four types"
		echo
		echo "  --init_extra:  use this switch to specify how non-core partitions are to"
		echo "be initialized. By default, non-core partitions are erased when created; however"
		echo "use this switch to initialize a partition, specifically a filesystem type"
		echo "For example, given a parts definition as follows"
		echo " --parts \"<core partitions>,1536k(extra1),2560k(extra2),-(data)\""  
		echo "could have a matching switch as follows:"
		echo " --init_extra \"extra1:yaffs2,extra2:raw,data:yaffs2\""
		echo "where the format is name:fstype, name must exactly match the partition "
		echo "name give to the --parts switch, and filesystem type which is yaffs or raw"
		echo "for a raw partition" 
		echo
		echo "NOTE:  This script is used to initialize an mtd flash device.  This script can"
		echo "       only be used of if the kernel modules have not already been loaded"
		echo
	fi

	exit
}

NAND=1
MTD_DEVNAME="NAND"
IMG_PREFIX="nand"
NOR=0
IMAGE_PATH=.
MTD_PARTS=
verbose=0
# don't initialize any non-core partitions
skip_extra=0
extra_parts_def=
# if extra partitions are defined on cmdline, flag it.
got_extra_defs=0
# image names
bootloader_img=
kernel_img=
initrd_img=
rootfs_img=
# default, we don't do rootfs on nor anyway
rootfs_type=yaffs2

if [ $# -eq 0 ]
then
	usage
fi

load_only=0
while [ $# -ne 0 ]
do
	if [ $1 == "--verbose" ]
	then
		verbose=1;
		shift;
	elif [ $1 == "--path" ]
	then
		IMAGE_PATH=$2
		shift;shift
	elif [ $1 == "--parts" ]
	then
		MTD_PARTS=$2
		shift;shift
	elif [ $1 == "--nand" ]
	then
		NAND=1
		NOR=0
		shift
	elif [ $1 == "--nor" ]
	then
		NOR=1
		NAND=0
		MTD_DEVNAME="NOR"
		IMG_PREFIX="nor"
		shift
	elif [ $1 == "--skip_extra" ]
	then
		skip_extra=1
		shift
	elif [ $1 == "--rootfs_type" ]
	then
		rootfs_type=$1
		shift;shift
	elif [ $1 == "--init_extra" ]
	then
		extra_parts_def=$2
		got_extra_defs=1
		shift;shift
	elif [ $1 == "--load" ]
	then
		load_only=1
		shift
	else
		usage
	fi
done

# the partition info is not updated frequently and may not be correct for the current release!!! 
if [ $NAND -eq  1 ]
then
	# as of 08/30/09
	DEFAULT_MTDPARTS="768k(mboot1),3200k(kernel1),3200k(kernel2),2176k(initrd1),2176k(initrd2),41472k(rootfs),-(data)"
elif [ $NOR -eq 1 ]
then
	# as of 08/30/09
	DEFAULT_MTDPARTS="1216k(kernel1),5888k(initrd1),9024k(data),256k(mboot1)"
fi

if [ $NAND -eq 1 ]
then
	MTD_MODULES="chipreg mtd mtdchar mtdblock"
	lsmod |egrep -sq '(aumb3000_nand|cmdlinepart)'
elif [ $NOR -eq 1 ]
then
	MTD_MODULES="chipreg mtd mtdchar mtdblock"
	lsmod |egrep -sq '(mg3500_flash|cmdlinepart)'
fi
if [ $? -eq 0 ]
then
	echo "ERROR:  Kernel modules already loaded! This script can only be run from an uninitialized state"
	exit
fi

if [ $load_only -eq 0 ]
then
	if [ "X${bootloader_img}" == "X" ]
	then
		bootloader_img=${IMG_PREFIX}-bootloader.bin
	fi
	if [ "X${kernel_img}" == "X" ]
	then
		kernel_img=${IMG_PREFIX}-kernel.img
	fi
	if [ "X${initrd_img}" == "X" ]
	then
		initrd_img=${IMG_PREFIX}-initrd.img
	fi
	if [ "X${rootfs_img}" == "X" -a $NAND -eq 1 ]
	then
	    # assume tgz, override w/ cmdline arg
		rootfs_img=${IMG_PREFIX}-rootfs.tgz
	fi

	if [ ! -e ${IMAGE_PATH}/${bootloader_img} ]
	then
		echo "ERROR: ${IMAGE_PATH}/${bootloader_img} not found"
		exit
	elif [ ! -e ${IMAGE_PATH}/${kernel_img} ]
	then
		echo "ERROR: ${IMAGE_PATH}/${kernel_img} not found"
		exit
	elif [ ! -e ${IMAGE_PATH}/${initrd_img} ]
	then
		echo "ERROR: ${IMAGE_PATH}/${initrd_img} not found"
		exit
	elif [ ! -e ${IMAGE_PATH}/${rootfs_img} ] && [ $NAND -eq 1 ]
	then
		echo "ERROR: ${IMAGE_PATH}/${rootfs_img} not found"
		exit
	fi
fi

if [ "X${MTD_PARTS}" == "X" ]
then
	echo "WARNING:  Using default ${MTD_DEVNAME} partition definition, this may not work!"
	echo "WARNING:  Refer to the current release notes to verify the following information"
	echo "${DEFAULT_MTDPARTS}"
	echo

	echo -ne "Do you want to continue (y/N) ? ";
	CHOICE=
	read CHOICE
	if [ "X${CHOICE}" == "Xy" ] || [ "X${CHOICE}" == "XY" ]
	then
		MTD_PARTS=${DEFAULT_MTDPARTS}
	else
		exit
	fi
fi

if [ $load_only -eq 0 ]
then
	echo -e "\nWARNING: All data on ${MTD_DEVNAME} 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
fi

load_modules()
{
	for module
	do
		load_module ${module}
		if [ $? -ne 0 ]
		then
			exit
		fi
	done
}

echo "Loading modules..."
load_modules ${MTD_MODULES}
if [ $NAND -eq 1 ]
then
	load_module cmdlinepart "mtdparts=aumb3000-nand:${MTD_PARTS}"
	if [ $? -ne 0 ]
	then
		exit
	fi
	load_modules nand aumb3000_nand
elif [ $NOR -eq 1 ]
then
	load_module cmdlinepart "mtdparts=physmap-flash:${MTD_PARTS}"
	if [ $? -ne 0 ]
	then
		exit
	fi
	load_modules mhif mg3500_flash
fi
# need some settle time
sleep 3

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

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

if [ $load_only -eq 1 ]
then
	echo "Init only, exiting..."
	exit
fi

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

	if [ $NOR -eq 1 ]
	then
		flashcp $filename $device
		if [ $? -ne 0 ]
		then
			echo "*****************************************************"
			echo "ERROR:  flashcp $filename $device failed... continuing"
			echo "*****************************************************"
		fi
	elif [ $NAND -eq 1 ]
	then
		nandwrite -p $device $filename
		if [ $? -ne 0 ]
		then
			echo "*******************************************************************"
			echo "ERROR:  nandwrite -p $device $filename failed... continuing"
			echo "*******************************************************************"
		fi
	fi
}

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

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
	else
		if [ $extra_start_index -eq 0 ]
		then 
			extra_start_index=$i
		fi
        type=extra
    fi

	if [ ! -e /dev/mtd${i} ]
	then
		echo "ERROR: /dev/mtd${i} does not exist, skipping $type initialization on this partition"
		continue
	fi

	if [ $type == "bootloader" ]
	then
		echo
		info "Initializing 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}/${bootloader_img}
	elif [ $type == "kernel" ]
	then
		echo
		info "Initializing 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}/${kernel_img}
	elif [ $type == "initrd" ]
	then
		echo
		info "Initializing 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}/${initrd_img}
	elif [ $type == "rootfs" ]
	then
		echo "${rootfs_img}" |grep -sq '\.bin'
		if [ $? -eq 0 ]
		then
			echo
			info "Initializing rootfs on /dev/mtdblock${i}"
			 # assume jffs2
			verbose "Erasing /dev/mtd${i}"
			flash_eraseall -j /dev/mtd${i}
			verbose "Writing /dev/mtd${i}"
			nandwrite -p /dev/mtd${i} ${IMAGE_PATH}/${rootfs_img}
		else
			mount | grep -sq /dev/mtdblock${i}
			if [ $? -eq 0 ]
			then
				echo "rootfs already mounted, can't do a live update!"
				exit
			fi
			# assume tgz, yaffs for now
			echo
			info "Initializing rootfs on /dev/mtdblock${i}"
			verbose "Erasing /dev/mtd${i}"
			flash_eraseall /dev/mtd${i}
			mkdir -p /rootfs_tmp
			info "Mounting ${rootfs_type} filesystem for rootfs"
			if [ ${rootfs_type} == "yaffs2" ]
			then
				FSTYPE=yaffs2
			fi
			mount -t ${FSTYPE} /dev/mtdblock${i} /rootfs_tmp
			if [ $? -ne 0 ]
			then
				echo "ERROR: mount -t ${FSTYPE} /dev/mtdblock${i} /rootfs_tmp failed... exiting"
				exit
			fi
			info "Untaring roofts..."
			if [ $verbose -eq 1 ]
			then
				tar -zxvf ${IMAGE_PATH}/${rootfs_img} -C /rootfs_tmp
			else
				tar -zxf ${IMAGE_PATH}/${rootfs_img} -C /rootfs_tmp
			fi
			if [ $? -ne 0 ]
			then
				echo "**************************************************************************"
				echo "ERROR: tar -zxf ${IMAGE_PATH}/${rootfs_img} -C /rootfs_tmp failed... exiting"
				echo "**************************************************************************"
				exit
			fi
			sync
			umount /rootfs_tmp
		fi
	elif [ $skip_extra -eq 0 ]
	then
	    # these are partitions with unrecognized partition names, these
        # may end up being raw partition or different kind of FS partitons
		echo
		info "Erasing spare partition for ${name} on  /dev/mtd${i}"
		flash_eraseall /dev/mtd${i}
	fi
	i=$((i+1));
done

if [ $skip_extra -eq 0 -a $got_extra_defs -eq 1 ]
then
	i=$extra_start_index
	for part in ${extra_parts_def}
	do
		# assume if not one of the "big 4" partitions, it will be an extra
	  	# partition which we will initialize as a yaffs filesystem but
		# but this does not mean it has to be used that way

		MOUNTPOINT=`echo ${part}|cut -d':' -f1`
		FSTYPE=`echo ${part}|cut -d':' -f2`

		if [ $NAND -eq 1 ]
		then
		    # yes, jffs2 could be added here too
			if [ "X${FSTYPE}" != "Xraw" -a "X${FSTYPE}" != "Xyaffs2" ]
			then
				echo "WARNING: Unsupported NAND partition type ${FSTYPE}, skipping /dev/mtd${i}"
				continue
			fi
		else
			if [ "X${FSTYPE}" != "Xraw" -a "X${FSTYPE}" != "Xjffs2" ]
			then
				echo "WARNING: Unsupported NOR partition type ${FSTYPE}, skipping /dev/mtd${i}"
				continue
			fi
		fi

		echo
		# skip raw partitions
		if [ "X${FSTYPE}" == "Xraw" ]
		then
			info "Raw partition /dev/mtd${i} initialized"
		else
			info "Creating ${FSTYPE} filesystem on /dev/mtdblock${i}"
			if [ ! -e /${MOUNTPOINT} ]
			then
				mkdir -p /${MOUNTPOINT}
			fi

			mount -t ${FSTYPE} /dev/mtdblock${i} /${MOUNTPOINT}
			if [ $? -ne 0 ]
			then
				echo "ERROR: mount -t ${FSTYPE} /dev/mtdblock${i} /${MOUNTPOINT} failed... continuing"
			else
				chmod 777 /${MOUNTPOINT}
				umount /${MOUNTPOINT}
				info "Partition /dev/mtd${i} initialized"
			fi
		fi

		i=$((i+1));
	done
fi

IFS=${OLD_IFS}
