#!/bin/bash

set -x

exec >> /var/log/generate-recovery-iso.log 2>&1

BASE="$(mktemp -d)"
AUFS="$(mktemp -d)"
TEMP="$(mktemp -d)"
UUID="$(uuidgen -r)"

if [ "$1" = "--oobe" ]; then
    mount --bind /cdrom "$BASE"
else
	RPLB=$(echo "get ubuntu-recovery/recovery_hotkey/partition_label" | debconf-communicate | cut -d" " -f 2)
	if [ -z "${RPLB}" ]
	then
		echo "recovery partition label not found"
		exit 1
	fi

	rpdev=$(blkid -L "${RPLB}")
	if [ -z "${rpdev}" ]; then
		rpdev=$(blkid -t PARTLABEL="Recovery Partition" -o device)
		if [ -z "${rpdev}" ]; then
			echo "recovery partition device not found"
			exit 1
		fi
	fi

    mount -v -t vfat -r "${rpdev}" "$BASE"
fi

#mount -v -t aufs -o br:$TEMP:$BASE none $AUFS
cp -rf "$BASE"/. "$AUFS"
mv -v "$AUFS"/syslinux "$AUFS"/isolinux
mv -v "$AUFS"/isolinux/syslinux.cfg "$AUFS"/isolinux/isolinux.cfg
mv -v "$AUFS"/boot/grub/grub.cfg.old "$AUFS"/boot/grub/grub.cfg
rm -f "$AUFS"/boot/grub/grubenv

echo "$UUID" > "$AUFS"/.disk/casper-uuid
INITRD="$(mktemp -d)"
cd "$INITRD" || exit
/usr/bin/unmkinitramfs "$AUFS"/casper/initrd .

# get compress method in initramfs
new_compress=$(grep "COMPRESS=" main/conf/initramfs.conf | awk -F"=" '{print $2}')
compress_command=""
if [ "$new_compress" == "gzip" ]; then
	compress_command="gzip -n"
elif [ "$new_compress" == "lzma" ] || [ "$new_compress" == "xz" ]; then
	compress_command="xz --check=crc32"
elif [ "$new_compress" == "lz4" ]; then
	compress_command="lz4 -9 -l"
elif [ "$new_compress" == "zstd" ]; then
    compress_command="zstd -q -19 -T0"
else
    echo "Unknown compress method ${new_compress}."
    exit 255
fi
#lzma -cd $AUFS/casper/initrd.lz -S lz | cpio -id
echo "$UUID" > main/conf/uuid.conf

:> "$AUFS"/casper/initrd
for component in 'early' 'early2' 'main';
do
	if [ $component == "main" ]; then
		(cd $component; find | cpio --quiet -o -H newc | $compress_command >> "$AUFS"/casper/initrd)
	else
		(cd $component; find | cpio --quiet -o -H newc  >> "$AUFS"/casper/initrd)
	fi
done	

#find | cpio --quiet -o -H newc | lzma -9c > $AUFS/casper/initrd.lz
cd - || exit
rm -fr "$INITRD"

if [ -d /sys/firmware/efi ]; then
    # efi mode
    # yuning: TODO: remove legacy support?

    if [ -f "$AUFS"/efi/boot/shimx64.efi ]; then
        mv -v "$AUFS"/efi/boot/shimx64.efi "$AUFS"/efi/boot/bootx64.efi
    fi

    if [ ! -d "$AUFS/efi" ]; then
        EFIFS="$(mktemp -d)"
        mount -v -t vfat "$AUFS"/boot/efi.img "$EFIFS"
        cp -R "$EFIFS"/efi "$AUFS"/
        umount "$EFIFS"
        rm -rf "$EFIFS"
    fi
else
    # legacy mode, remove efi support
    rm -rf "$AUFS"/boot/grub
    rm -rf "$AUFS"/boot/efi.img
    rm -rf "$AUFS"/efi/
fi

mkdir -p /usr/share/ubuntu/

if [ -e "$AUFS"/boot/efi.img ]
then
    UEFI_OPTS="-eltorito-alt-boot -e boot/efi.img -no-emul-boot"
else
    UEFI_OPTS=""
fi

# We have preserved one efi.img under /boot dir, we can use it to 
# create UEFI/Legacy hybrid ISO.
genisoimage \
    -quiet \
    -o /usr/share/ubuntu/recovery.tmp \
    -input-charset utf-8 \
    -boot-load-size 4 \
    -boot-info-table \
    -pad \
    -r \
    -J \
    -joliet-long \
    -N \
    -hide-joliet-trans-tbl \
    -cache-inodes \
    -l \
    -publisher 'Canonical Ltd.' \
    -sysid 'Ubuntu' \
	-V 'Reinstallation Media' \
	-volset "$(tail -1 /etc/buildstamp)" \
	-A 'Ubuntu Recovery Media Creator' \
	-b isolinux/isolinux.bin \
	-c isolinux/boot.cat \
	-no-emul-boot $UEFI_OPTS "$AUFS"
rm -rf "$AUFS"
umount "$BASE"
rmdir "$BASE"
find "$TEMP"
rm -fr "$TEMP"
mv /usr/share/ubuntu/recovery.tmp /usr/share/ubuntu/recovery.iso
if [ -x /usr/bin/isohybrid ]
then
    isohybrid --uefi /usr/share/ubuntu/recovery.iso
fi
