#!/bin/bash
# Requires: qemu, ovmf, qemu-system-x86, please enable the kvm at BIOS.
# Usage:
# ./sutton-factory-image-creator-uefi.sh sutton-xxx-gm.iso
#
# After that, copy the dd image to USB stick, boot up with USB and at the install screen
# 'Ctrl + Alt + F2' to switch tty2.
# mkdir usb 
# mount /dev/sda2 usb (because the hdd device on target machine is /dev/nvme0n1*, so the USB stick will be "/dev/sda")
# or
# mount /dev/sdb2 usb (because the hdd device on target machine is /dev/sda*, so the USB stick will be "/dev/sdb")
# cd usb
# gzip -dc sutton-nvme.img.dd.gz | dd of=/dev/nvme0n1
# or
# gzip -dc sutton-sata.img.dd.gz | dd of=/dev/sda
#
# Maintainer: shengyao.xue@canonical.com or bin.li@canonical.com

#set -x
set -e

if [ "$(id -u)" != "0" ]; then
	echo "Must be run as root"
	exit 1
fi

if [ $# != 1 ]; then
	cat <<EOF
Usage: $0 sutton.iso
EOF
	exit 1
fi

# detect how to compress
if type pigz >/dev/null 2>&1; then
	compress="pigz"
elif type gzip >/dev/null 2>&1; then
	compress="gzip"
else
	echo "abort: pigz or gzip not found"
	echo "tips: apt-get install pigz"
	exit 1
fi

qemu="qemu-system-x86_64"
if ! type "$qemu" >/dev/null 2>&1; then
	echo "abort: $qemu not found"
	echo "tips: apt-get install qemu-system-x86"
	exit 1
fi

if ! type "qemu-img" >/dev/null 2>&1; then
	echo "abort: qemu-img not found"
	echo "tips: apt-get install qemu-utils"
	exit 1
fi

# install the iso

iso="$1"
if [ ! -e "$iso" ]; then
	echo "abort: file not found: $iso"
	exit 1
fi

extname=${iso##*.}
version=$(basename "$iso" ".$extname")

opts=""
opts="$opts -m 2G"
opts="$opts -machine accel=kvm:xen:tcg"
opts="$opts -net none"
opts="$opts -no-reboot"
opts="$opts -bios OVMF.fd"

echo "The storage device of the target machine is nvme or sata? Type 'nvme' or 'sata' and press [ENTER]:"
read nors

if [ "$nors" == "nvme" ]; then
	hdd="${version}-${nors}.img"
	opts="$opts -drive file=$hdd,if=none,id=drv0 -device nvme,drive=drv0,serial=foo"
elif [ "$nors" == "sata" ]; then
	hdd="${version}-${nors}.img"
	opts="$opts -drive file=$hdd"
else
	echo "storage type (nvme or sata) input error."
	exit 1
fi

ddimg="${hdd}.dd.gz"

size=32000
# must be larger than $size
# the hdd_size need bigger than 64G
hdd_size=60000

# TODO: check for KVM and XEN

rm -f "$hdd" "$ddimg"

# create hdd image
echo "create virtual hdd image"
qemu-img create -f qcow2 "$hdd" "${hdd_size}M"


# P1
echo "install stage 1: create local recovery partition"
"$qemu" $opts -cdrom "$iso"
# P2
echo "install stage 2: recovery system"
"$qemu" $opts

#echo "Done"
#exit

if lsmod | grep -q nbd; then
	qemu-nbd -d /dev/nbd15 || true
	rmmod nbd
fi
modprobe nbd

qemu-nbd -c /dev/nbd15 $hdd

if [ ! -d mnt ]; then
	mkdir mnt
fi
udevadm settle
mount /dev/nbd15p1 mnt
rm -rf mnt/EFI/BOOT
cp -r mnt/EFI/ubuntu mnt/EFI/BOOT
udevadm settle
umount mnt

mount /dev/nbd15p3 mnt
sed -i 's/python <<EOF/python3 <<EOF/g' mnt/usr/share/volatile/task/00_partition
sed -i 's/print p.before/print(p.before)/g' mnt/usr/share/volatile/task/00_partition
udevadm settle
umount mnt
rmdir mnt

# create dd image
echo "create dd image, this may take several minutes..."
#dd if="$hdd" bs=1M
dd if="/dev/nbd15" bs=1M count="$size" | $compress -9c > "$ddimg"

# calculate checksum
for method in md5sum sha1sum; do
	"$method" "$ddimg" > "${ddimg}.${method}"
done

qemu-nbd -d /dev/nbd15
rm -f "$hdd"

echo "Finished!"
