#!/usr/bin/env bash

# configuration variables
MOUNT=/mnt/rootfs
DISKS=(/dev/sda /dev/sdb)
RAIDVOL=/dev/md0
IMAGE_PATH=/livemnt/boot
IMAGE_PATH_ALT=/mnt/usb
ROOT_MIN_SIZE=2048

# globals
cpu=
memory=

function print() {
	echo -e "\033[33m$1\033[0m"
}

function die() {
	echo -e "\033[31m$1\033[0m"
	exit 1
}

if ls ${IMAGE_PATH}/image_dosis_2_*.tar.bz2 > /dev/null 2>&1; then
    IMAGE=`ls ${IMAGE_PATH}/image_dosis_2_*.tar.bz2 | sort -rV | head -1`
elif ls ${IMAGE_PATH_ALT}/image_dosis_2_*.tar.bz2 > /dev/null 2>&1; then
    IMAGE=`ls ${IMAGE_PATH_ALT}/image_dosis_2_*.tar.bz2 | sort -rV | head -1`
else
    die "Unable to find an image!"
fi
print "Image found: ${IMAGE}"

print "Stopping RAID"
mdadm --zero-superblock ${RAIDVOL}
dmsetup remove_all
mdadm --stop --scan

print "Lobotomize disks"
for disk in ${DISKS[@]}; do
    dd if=/dev/zero of=${disk} bs=1M count=300
    wipefs --all --force ${disk}
done

sleep 3

for disk in ${DISKS[@]}; do
	print "Probing ${disk}"
	partprobe "${disk}"
done

function sysstat() {
	cpu=`cat /proc/cpuinfo | grep -i "model name" | head -1 | cut -d: -f2 | sed 's/[[:space:]]*$//' | sed 's/^[[:space:]]*//'`
    memory=`free -m | grep Mem: | awk '{print $2}'`
    print "System Information"
	print "------------------"
	print "CPU: ${cpu}"
	print "MEM: ${memory}MB"
	for disk in ${DISKS[@]}; do
	    print "${disk}: `echo "$(blockdev --getsize64 ${disk})/1024/1024" | bc`MB"
    done
	print
}

function efi_partition_configuration() {
	echo "2048,409600,,*"
	echo ",,,"
}

function legacy_partition_configuration() {
	echo "2048,,,*"
}

function booted_efi_mode() {
	if [ -d /sys/firmware/efi ]; then
		return 0
	else
		return 1
	fi
}

function setupraid() {
	if booted_efi_mode; then
		for disk in ${DISKS[@]}; do
			print "Partitioning ${disk}"
			efi_partition_configuration | sfdisk --force -q ${disk}
		done

		print "Creating RAID volume ${RAIDVOL}"
		mdadm --create ${RAIDVOL} \
			--level=1 \
			--raid-devices=2 \
			--metadata=1.2 \
			`for disk in ${DISKS[@]}; do echo ${disk}2; done`

		sleep 3

		for disk in ${DISKS[@]}; do
			yes | mkfs.vfat ${disk}1
		done
	else
		for disk in ${DISKS[@]}; do
			print "Partitioning ${disk}"
			legacy_partition_configuration | sfdisk --force -q ${disk}
		done

		print "Creating RAID volume ${RAIDVOL}"
		mdadm --create ${RAIDVOL} \
			--level=1 \
			--raid-devices=2 \
			--metadata=1.2 \
			`for disk in ${DISKS[@]}; do echo ${disk}1; done`

		sleep 3
	fi

	print "Formatting ${RAIDVOL}"
	yes | mkfs.ext4 ${RAIDVOL} || die "Failed to partition disk"
}

function check_rootfs() {
	print "Checking for Rootfs ${IMAGE}"
	print "----------------------------"
	[ -f "${IMAGE}" ] || die "${IMAGE} does not exist"
}

function umount_filesystems() {
	umount "${MOUNT}/sys"
	umount "${MOUNT}/dev"
	umount "${MOUNT}/proc"
	umount "${MOUNT}/boot/efi"
	umount "${RAIDVOL}"
}
trap umount_filesystems EXIT

function mount_filesystems() {
	print "Mounting Filesystems"
	print "--------------------"
	mkdir -p "${MOUNT}"
	mount "${RAIDVOL}" "${MOUNT}"

	if booted_efi_mode; then
		mkdir -p "${MOUNT}/boot/efi"
		mount ${DISKS[0]}1 "${MOUNT}/boot/efi"
	fi
}

function copy_efi_contents() {
	if booted_efi_mode; then
		print "Duplicating EFI contents on mirrored disk"
		disk1="`mktemp -d`"
		disk2="`mktemp -d`"
		mount ${DISKS[0]}1 "${disk1}"
		mount ${DISKS[1]}1 "${disk2}"
		cp -rpv "${disk1}"/* "${disk2}"/.
		umount "${disk1}"
		umount "${disk2}"
		rmdir "${disk1}"
		rmdir "${disk2}"
	fi
}

function extract_rootfs() {
	print "Extracting Rootfs Image"
	print "-----------------------"
	pv "${IMAGE}" | tar xpjf - -C "${MOUNT}" || die "failed to extract ${IMAGE}"
}

function create_swapfile() {
    print "Creating /swapfile"
	chroot "${MOUNT}" /usr/share/dosis2/bin/setup-swap
}

function mount_extra_filesystems() {
	mount --bind /dev "${MOUNT}/dev"
	mount -t proc none "${MOUNT}/proc"
	mount -t sysfs none "${MOUNT}/sys"
}

function get_uuid() {
	lsblk -n -o UUID "$1"
}

function gen_fstab() {
	root_uuid=`get_uuid "${RAIDVOL}"`
    lineformat="%-45s %-10s %-7s %-15s %-3s\n"
    echo "### This file was generated by install.sh"
    echo "### You may modify its contents."
    echo "###"
    echo ""
    printf "${lineformat}" "UUID=${root_uuid}" "/" "ext4" "defaults" "1 1"
    printf "${lineformat}" "/swapfile" "swap" "swap" "defaults" "0 0"
}

function write_fstab() {
	fstab_="${MOUNT}/etc/fstab"

    cat <<EOF > "${fstab_}"
`gen_fstab`
EOF

	print "Current fstab"
	print "-------------"
	cat "${fstab_}"
}

function cleanup_authkeys() {
	rm -f "${MOUNT}/root/.ssh/authorized_keys"
}

function cleanup_mnt() {
	rm -rf "${MOUNT}/mnt/*"
}

function setup_grub() {
	print "Installing and Updating Grub"
	print "------------------------"
	for disk in ${DISKS[@]}; do
	    chroot "${MOUNT}" /usr/sbin/grub-install "${disk}"
    done
	chroot "${MOUNT}" /usr/sbin/update-grub
}

function update_initrd() {
	print "Updating Initial RAM Filesystem"
	print "-------------------------------"

	chroot "${MOUNT}" /usr/sbin/update-initramfs -u
}

sysstat
setupraid
check_rootfs
mount_filesystems
extract_rootfs
cleanup_mnt
mount_extra_filesystems
create_swapfile
write_fstab
cleanup_authkeys
setup_grub
update_initrd
umount_filesystems
copy_efi_contents

reboot
