#! /bin/bash
#

set -e
set -u

ProgName="$0"
KVer=''
KV1=''
KV2=''

ERROR=0
Arch=$(uname -m)
WGET='/usr/bin/wget'
CURL='/usr/bin/curl'
MKIMAGE='/usr/bin/mkimage'
DEPMOD='/usr/sbin/depmod'
[[ -f $DEPMOD ]] || DEPMOD='/sbin/depmod'
FLASH_ERASEALL='/usr/sbin/flash_eraseall'
FLASH_ERASE='/usr/sbin/flash_erase'
NANDWRITE='/usr/sbin/nandwrite'
SITE1="http://www.xilka.com/kernel"
IS_DTB='1'

AssumedDevice=''
DTBS=''

CHECK_PROGRAMS="$WGET $CURL $DEPMOD"

if [[ $Arch == aarch64 ]]; then
	Platform='arm64'
	AssumedDevice='unknown'
	model="$(cat /proc/device-tree/model 2>/dev/null)"

	if [[ $model == *'RockPro64'* ]]; then
		Platform='rock64'
		AssumedDevice='rockpro64'
	elif [[ $model == *'Khadas VIM2'* ]]; then
		Platform='meson'
		AssumedDevice='vim2'
	elif [[ $model == *'Raspberry Pi 5'* ]]; then
		Platform='pi564'
		AssumedDevice='pi5'
	elif [[ $model == *'Raspberry Pi 4'* ]] \
			|| [[ $model == *'Raspberry Pi Compute Module 4'* ]]; then
		Platform='pi464'
		AssumedDevice='pi4'
	elif [[ $model == *'Raspberry Pi 3'* ]] \
			|| [[ $model == *'Raspberry Pi Compute Module 3'* ]] \
			|| [[ $model == *'Raspberry Pi Zero 2'* ]]; then
		Platform='pi8'
		AssumedDevice='pi3'
	elif [[ $model == *'ESPRESSOBin'* ]] \
			|| [[ $model == *'3720 Community'* ]]; then
		AssumedDevice='espressobin'
	fi
elif [[ $Arch == armv7l ]]; then
	Platform='dove'
	model="$(cat /proc/device-tree/model 2>/dev/null)"

	if grep -qs 'i.MX6' /proc/cpuinfo \
		|| grep -qs 'SolidRun i.MX' /proc/cpuinfo \
		|| grep -qs 'SolidRun Cubox-i' /proc/cpuinfo; then
		Platform='imx'
		AssumedDevice='imx'
	elif grep -qs 'CuBox' /proc/cpuinfo || grep -qs 'CuBox' /proc/device-tree/model; then
		AssumedDevice='cubox'
		CHECK_PROGRAMS="$CHECK_PROGRAMS $MKIMAGE"
	elif [[ $model == *'Raspberry Pi 4'* ]] \
			|| [[ $model == *'Raspberry Pi Compute Module 4'* ]]; then
		Platform='pi432'
		AssumedDevice='pi4'
	elif [[ $model == *'Raspberry Pi 3'* ]] \
		|| [[ $model == *'Raspberry Pi Compute Module 3'* ]] \
		|| [[ $model == *'Raspberry Pi Zero 2'* ]]; then
		Platform='pi7'
		AssumedDevice='pi3'
	elif [[ $model == *'Raspberry Pi 2'* ]]; then
		Platform='pi7'
		AssumedDevice='pi2'
	fi
elif [[ $Arch == armv6l ]] \
		&& (grep -qs 'BCM2708' '/proc/cpuinfo' \
			|| grep -qs 'BCM2835' '/proc/cpuinfo'); then
	Platform='pi'
	AssumedDevice='pi'
elif [[ $Arch == armv5tel ]]; then
	Platform='kirkwood'
	if [[ -f /proc/device-tree/model ]]; then
		model="$(cat /proc/device-tree/model 2>/dev/null)"
	else
		model=$(grep "^Hardware" /proc/cpuinfo | sed -re "s/.* ([^ ]*Plug|Tonido2)($| .*)/\1/")
		[[ $model =~ 'GuruPlug' && ! -c /dev/mtd0 ]] && model='DreamPlug'
		[[ $model =~ 'SheevaPlug' && -d /sys/devices/platform/sata_mv.0 ]] && model="eSATA $model"
		[[ $model =~ 'Tonido2' ]] && model='Topkick'
	fi
	model=${model/Plug/plug}
	if   [[ $model =~ 'Dreamplug' ]]; then
		AssumedDevice='dreamplug'
	elif [[ $model =~ 'Guruplug' ]]; then
		AssumedDevice='guruplug-server-plus'
	elif [[ $model =~ 'Sheevaplug' ]]; then
		AssumedDevice='sheevaplug'
		[[ ! $model =~ 'eSATA' ]] || AssumedDevice+='-esata'
	elif [[ $model =~ 'Topkick' ]]; then
		AssumedDevice='topkick'
	fi

	if [[ $AssumedDevice == 'dreamplug' ]]; then
		CHECK_PROGRAMS="$CHECK_PROGRAMS $MKIMAGE"
	elif [[ -n $AssumedDevice ]]; then
		CHECK_PROGRAMS="$CHECK_PROGRAMS $MKIMAGE $FLASH_ERASEALL $FLASH_ERASE"
	fi
else
	"Unknown architecture: $Arch"
	exit 1
fi

echo "Assuming we're on a $Platform-$AssumedDevice"

for f in $CHECK_PROGRAMS; do
	if [[ ! -x $f ]]; then
		echo "missing $f"
		ERROR=1
	fi
done

if (( ERROR )); then
	exit 1
fi

function Usage()
{
	if [[ -z $AssumedDevice ]]; then
		echo "Use:  $ProgName VERSION DEVICE to write kernel to /boot"
		echo "          where VERSION is something like 4.19.13"
		echo "          and DEVICE is something like cubox or sheevaplug"
	else
		echo "Use:  $ProgName VERSION to write kernel to /boot"
		echo "          where VERSION is something like 4.19.13"
	fi
}

function Sha1Compare ()
{
	if [[ $(cat $2 | cut -d' ' -f1) \
			!= $(sha1sum $1 | cut -d' ' -f1) ]]; then
		echo "Bad sha1 detected on $1"
		exit 1
	fi
}

function GetDownloadPath()
{
	local site=$1

	local v1=$(echo $KVer | cut -d'.' -f1)
	local v2=$(echo $KVer | cut -d'.' -f2)
	local v3=$(echo $KVer | cut -d'.' -f3)
	local v4=$(echo $KVer | cut -d'.' -f4)
	local path="$v1/$v1.$v2"

	if (( $v1 == 2 )); then
		if [[ -z $v3 ]]; then
			path="$path/$v1.$v2"
		else
			path="$path/$v1.$v2.$v3"
			if [[ -z $v4 ]]; then
				path="$path/$v1.$v2.$v3"
			else
				path="$path/$v1.$v2.$v3.$v4"
			fi
		fi
	else
		if [[ -z $v3 ]]; then
			path="$path/$v1.$v2"
		elif [[ ! -z $v4 ]]; then
			path="$path/$v1.$v2.$v3/$v1.$v2.$v3.$v4"
		else
			path="$path/$v1.$v2.$v3"
		fi
	fi

	for (( i = 5; i >= 1; i-- )); do
		testpath=$site/$path/release/$i
		if $CURL -sf $testpath; then
			path=$path/release/$i
			break
		fi
	done

	echo $path
}

function DownloadOld()
{
	echo "Downloading files"

	local path=''
	local f=''
	local pfx=''
	local ext=''
	local file=''
	local image='zImage'

	if [[ $Arch == 'aarch64' ]]; then
		image='Image'
	elif [[ $KVer == 3.5.7.* ]]; then
		IS_DTB='0'
		image='uImage'
	fi

	for f in $image Modules.tar.gz System.map; do
		pfx="$AssumedDevice-$KVer-$f"

		if (( IS_DTB )); then
			pfx="$Platform-$KVer-$f"
		fi

		for ext in '' '.sha1'; do
			file="$pfx$ext"

			if [[ ! -f $file ]]; then
				[[ ! -z $path ]] || path=$(GetDownloadPath $SITE1)

				local s=''
				for s in $SITE1; do
					if $WGET -c $s/$path/$file && [[ -f $file ]]; then
						break
					fi
				done
				if [[ ! -f $file ]]; then
					echo "Unable to retrieve $file"
					exit 1
				fi
			fi
		done

		Sha1Compare $pfx{,'.sha1'}
	done

	if (( IS_DTB )); then
		local dtb=''

		for dtb in $DTBS; do
			for ext in '' '.sha1'; do
				if [[ ! -f $dtb$ext ]]; then
					local s=''
					for s in $SITE1; do
						if $WGET -c $s/$path/$dtb$ext && [[ -f $dtb$ext ]]; then
							break
						fi
					done
					if [[ ! -f $dtb$ext ]]; then
						echo "Unable to retrieve $dtb$ext"
						exit 1
					fi
				fi
			done

			Sha1Compare $dtb{,'.sha1'}
		done
	fi
}

function DownloadCommon()
{
	local device="$1"
	echo "Downloading files"

	local path=''
	local f=''
	local imageExt='-Image'
	[[ $device != pi* ]] || imageExt='.img'
	local ext=''

	for ext in '-Modules.tar.gz' '-Modules.tar.gz.sha1' $imageExt ${imageExt}.sha1 \
			'-dtbs.tar.gz' '-dtbs.tar.gz.sha1' '-System.map' '-System.map.sha1'; do
		if [[ ! -f $device-$KVer$ext ]]; then
			[[ ! -z $path ]] || path=$(GetDownloadPath $SITE1)

			local s=''
			for s in $SITE1; do
				if $WGET -c $s/$path/$device-$KVer$ext && [[ -f $device-$KVer$ext ]]; then
					break
				fi
			done
			if [[ ! -f $device-$KVer$ext ]]; then
				echo "Unable to retrieve $device-$KVer$ext"
				exit 1
			fi
		fi
	done

	Sha1Compare $device-$KVer-dtbs.tar.gz{,.sha1}
	Sha1Compare $device-$KVer-Modules.tar.gz{,.sha1}
	Sha1Compare $device-$KVer$imageExt{,.sha1}
	Sha1Compare $device-$KVer-System.map{,.sha1}
}

MkImage()
{
	cat $Platform-$KVer-zImage $Platform-$Device-$KVer.dtb > /tmp/zImage.$Device$$

	$MKIMAGE -A arm -O linux -C none  -T kernel -a 0x00008000 -e 0x00008000 \
		-n "Linux-$Platform-$Device-$KVer" \
		-d /tmp/zImage.$Device$$ \
		$Platform-$Device-$KVer-uImage

	rm /tmp/zImage.$Device$$
}

function ExtractModulesOld()
{
	local v1=$(echo $KVer | cut -d'.' -f1)
	local v2=$(echo $KVer | cut -d'.' -f2)
	local v3=$(echo $KVer | cut -d'.' -f3)
	local depmodVer=$KVer
	[[ ! -z $v3 ]] || depmodVer="$v1.$v2.0"
	local extractDir='/'
	[[ ! -L '/lib' ]] || extractDir=$(dirname $(readlink '/lib'))
	[[ -d $extractDir ]] || extractDir="/$extractDir"

	echo "Extracting modules"
	local arc=$AssumedDevice-$KVer-Modules.tar.gz
	local map=$AssumedDevice-$KVer-System.map

	if (( IS_DTB )); then
		arc=$Platform-$KVer-Modules.tar.gz
		map=$Platform-$KVer-System.map
	fi

	tar x -C $extractDir --overwrite -zf $arc
	[[ -d /boot ]] || mkdir /boot
	cp $map /boot/
	$DEPMOD -eF /boot/$map $depmodVer
}

function ExtractModulesCommon()
{
	local device="$1"
	local v1=$(echo $KVer | cut -d'.' -f1)
	local v2=$(echo $KVer | cut -d'.' -f2)
	local v3=$(echo $KVer | cut -d'.' -f3)
	local depmodVer=$KVer
	[[ $device != 'pi564' ]] || depmodVer="$KVer-v8-16k"
	[[ $device != 'pi464' ]] || depmodVer="$KVer-v8"
	[[ $device != 'pi432' ]] || depmodVer="$KVer-v7l"
	[[ $device != 'pi8' ]] || depmodVer="$KVer-v8"
	[[ $device != 'pi7l' ]] || depmodVer="$KVer-v7l"
	[[ $device != 'pi7' ]] || depmodVer="$KVer-v7"

	[[ ! -z $v3 ]] || depmodVer="$v1.$v2.0"
	[[ ! -d /lib/modules/${depmodVer}+ ]] || depmodVer="${depmodVer}+"

	local extractDir='/'
	[[ ! -L '/lib' ]] || extractDir=$(dirname $(readlink '/lib'))
	[[ -d $extractDir ]] || extractDir="/$extractDir"

	echo "Extracting modules"
	tar x -C $extractDir --overwrite -zf $device-$KVer-Modules.tar.gz
	[[ -d /boot ]] || mkdir /boot
	cp $device-$KVer-dtbs.tar.gz /boot/
	cp $device-$KVer-System.map /boot/

	$DEPMOD -eF /boot/$device-$KVer-System.map $depmodVer
}

function CheckSize()
{
	local uImageFile=$Platform-$Device-$KVer-uImage
	local mtd=$(grep uImage /proc/mtd | cut -d' ' -f1 | sed 's#:##')
	local mtdHexSize=$(grep uImage /proc/mtd | cut -d' ' -f2)
	local mtdSize=''
	let mtdSize=0x$mtdHexSize
	local uImageSize=$(stat -c%s "$uImageFile")
	if (( uImageSize > mtdSize )); then
		echo "uImage size $uImageSize exceeds $mtd size $mtdSize"
		echo "You must resize your mtds before this kernel will fit!"
		echo "The simplest way to do this is with the SheevaPlug installer"
		exit 1
	fi
}

function FlashEraseAll()
{
	if [[ -z $FLASH_ERASEALL ]] || file $FLASH_ERASEALL | grep -qs 'shell script'; then
		$FLASH_ERASE $@ 0 0
	else
		$FLASH_ERASEALL $@
	fi
}

function NandKernel()
{
	echo "Flashing kernel to NAND"
	# MTDDEVICE should be either /dev/mtd0 or /dev/mtd1
	# depending on which partition contains the kernel uImage

	local Mtd=$(grep uImage /proc/mtd | cut -d' ' -f1 | sed 's#:##')

	FlashEraseAll -j /dev/$Mtd
	$NANDWRITE -pm /dev/$Mtd $Platform-$Device-$KVer-uImage
}

function RootKernelOld()
{
	echo "Writing kernel to /boot"
	[[ -d /boot ]] || mkdir /boot

	local img="$AssumedDevice-$KVer-uImage"

	if (( IS_DTB )); then
		if [[ $Platform == 'imx' ]]; then
			img="$Platform-$KVer-zImage"
		else
			img="$Platform-$Device-$KVer-uImage"
		fi
	fi

	if [[ $Platform == 'imx' ]]; then
		cp $img $DTBS /boot/
		echo "****************************************************************"
		echo " now run sudo ./SET-IMX-KERNEL.sh $KVer in /boot"
		echo "****************************************************************"
	else
		cp $img /boot/
		echo "****************************************************************"
		echo " create a link from /boot/$img to /boot/uImage"
		echo "****************************************************************"
	fi
}

function RootKernelCommon()
{
	local scriptName="SET-ARM64-KERNEL.sh"
	local imageExt='-Image'

	if [[ $AssumedDevice == vim2 ]]; then
		scriptName="SET-VIM2-KERNEL.sh"
	elif [[ $Platform == pi* ]]; then
		scriptName="SET-PI-KERNEL.sh"
		imageExt='.img'
	fi

	echo "Writing kernel to /boot"
	[[ -d /boot ]] || mkdir /boot
	cp $Platform-$KVer$imageExt /boot/
	echo "****************************************************************"
	echo " now run sudo ./$scriptName $KVer in /boot"
	echo "****************************************************************"
}

if (( $# < 1 )) || [[ $1 == "--help" ]]; then
	Usage
	exit 1
elif (( $# < 2 )) && [[ -z $AssumedDevice ]]; then
	Usage
	exit 1
fi

KVer="$1"
KV1=$(echo $KVer | cut -d'.' -f1)
KV2=$(echo $KVer | cut -d'.' -f2)

if (( $# >= 2 )); then
	Device="$2"
	echo "Overriding to $Platform-$Device"
else
	Device="$AssumedDevice"
fi

DTBS="$Platform-$Device-$KVer.dtb"

if [[ $Platform == imx ]]; then
	DTBS="imx6dl-cubox-i-$KVer.dtb  imx6q-cubox-i-$KVer.dtb"
	DTBS="$DTBS imx6dl-hummingboard-$KVer.dtb imx6q-hummingboard-$KVer.dtb"

	if (( KV1 == 4 && KV2 >= 4 )) || (( KV1 > 4 )); then
		DTBS="$DTBS imx6dl-hummingboard2-$KVer.dtb imx6q-hummingboard2-$KVer.dtb"
	fi
fi

if [[ $Arch == aarch64 ]] || [[ $Platform == pi* ]]; then
	DownloadCommon $Platform
	ExtractModulesCommon $Platform
	RootKernelCommon $Platform
else
	DoNandKernel=0

	if [[ $Platform == 'kirkwood' ]] && [[ $AssumedDevice != 'dreamplug' ]]; then
		echo "n - write uImage to nand"
		echo "r - write uImage to /root"
		echo -n "? "
		read response
		if [[ $response == 'n' ]]; then
			DoNandKernel=1
		elif [[ $response != 'r' ]]; then
			exit 1
		fi
	fi

	DownloadOld

	if (( IS_DTB )) && [[ $Platform != 'imx' ]]; then
		MkImage
	fi

	ExtractModulesOld

	if (( DoNandKernel )); then
		CheckSize
		NandKernel
	else
		RootKernelOld
	fi
fi