#!/bin/bash

set -eu

ISO_MNT="/mnt/iso"
PARTITION_LABEL="multiboot"
MULTIBOOT_MNT="/mnt/multiboot"
SYSLINUX_VERSION="6.03"
SYSLINUX_URL="https://www.kernel.org/pub/linux/utils/boot/syslinux/syslinux-${SYSLINUX_VERSION}.tar.gz"

function cmd_format {
    if [ ! -b "${1:-}" ]; then
        echo "Usage: $0 format /dev/sdX"
        exit 1
    fi
    set -x

    sudo -k
    DISK_DEVICE="$1"
    PARTITION_DEVICE="${DISK_DEVICE}1"
    echo -ne 'label: dos\ntype=c, bootable\n' | sudo sfdisk "$DISK_DEVICE"
    sudo mkfs.vfat -n "$PARTITION_LABEL" "$PARTITION_DEVICE"
}

function cmd_bootloader {
    DISK_DEVICE="$(mount|grep /mnt/multiboot|cut -d' ' -f1|sed 's/[0-9]*$//')"
    if [ ! -b "$DISK_DEVICE" ]; then
        echo "ERROR! Could not find disk to install bootloader. Try using mount."
        exit 1
    fi
    set -x

    sudo -k
    install_grub_bios
    install_grub_efi
}

function install_grub_bios {
    sudo grub-install \
        --target=i386-pc \
        --boot-directory="$MULTIBOOT_MNT" \
        "$DISK_DEVICE"
}

function install_grub_efi {
    sudo grub-install \
        --target=x86_64-efi \
        --no-nvram \
        --removable \
        --efi-directory="$MULTIBOOT_MNT" \
        --boot-directory="$MULTIBOOT_MNT" \
        "$DISK_DEVICE"
}

function cmd_mount {
    set -x

    PARTITION_DEVICE="$(readlink -f /dev/disk/by-label/multiboot)"
    sudo mkdir -p "$MULTIBOOT_MNT"
    while sudo umount "$PARTITION_DEVICE" &> /dev/null; do true; done
    sudo mount "$PARTITION_DEVICE" "$MULTIBOOT_MNT" -o "uid=$(whoami)"
}

function cmd_umount {
    set -x

    sudo umount "$MULTIBOOT_MNT"
    sudo rmdir "$MULTIBOOT_MNT"
}

function cmd_add_iso {
    if [ ! -f "${1:-}" ]; then
        echo "Usage: $0 add_iso /path/to/ubuntu.iso"
        exit 1
    fi
    set -x

    ISO_FILENAME="$1"
    ISO_BASENAME="$(basename "$ISO_FILENAME")"
    GRUB_PART_CFG="${MULTIBOOT_MNT}/grub/parts/${ISO_BASENAME}"
    local f="$(basename "$ISO_FILENAME")"

    mkdir -p "$(dirname "$GRUB_PART_CFG")"

    if [[ "$ISO_BASENAME" == ubuntu-*-desktop-*.iso ]]; then
        setup_ubuntu
    elif [[ "$ISO_BASENAME" == debian-live-*.iso ]]; then
        setup_debian
    elif [[ "$ISO_BASENAME" == Fedora-Workstation-Live-*.iso ]]; then
        setup_fedora
    elif [[ "$ISO_BASENAME" == archlinux-*.iso ]]; then
        setup_archlinux
    elif [[ "$ISO_BASENAME" == FD12*.zip ]]; then
        setup_freedos
    else
        echo "Unsupported ISO! $ISO_BASENAME"
    fi

    delete_stray_grub_parts
    print_grub_cfg > "${MULTIBOOT_MNT}/grub/grub.cfg"
}

function setup_ubuntu {
    copy_iso_data

    cat <<EOT > "$GRUB_PART_CFG"
menuentry '$ISO_BASENAME' {
  loopback loop /$ISO_BASENAME
  linux (loop)/casper/vmlinuz.efi boot=casper quiet iso-scan/filename=/$ISO_BASENAME
  initrd (loop)/casper/initrd.lz
}

EOT
}

function setup_debian {
    copy_iso_data

    cat <<EOT > "$GRUB_PART_CFG"
menuentry '$ISO_BASENAME' {
  loopback loop /$ISO_BASENAME
  linux (loop)/live/vmlinuz boot=live components findiso=/$ISO_BASENAME
  initrd (loop)/live/initrd.img
}

EOT
}

function setup_fedora {
    copy_iso_data

    cat <<EOT > "$GRUB_PART_CFG"
menuentry '$ISO_BASENAME' {
  loopback loop /$ISO_BASENAME
  linux (loop)/isolinux/vmlinuz root=live:CDLABEL=$(blkid -s LABEL -o value "$ISO_FILENAME") rd.live.image quiet iso-scan/filename=/$ISO_BASENAME
  initrd (loop)/isolinux/initrd.img
}

EOT
}

function setup_archlinux {
    copy_iso_data

    cat <<EOT > "$GRUB_PART_CFG"
menuentry '$ISO_BASENAME' {
  loopback loop /$ISO_BASENAME
  linux (loop)/arch/boot/x86_64/vmlinuz img_label=${PARTITION_LABEL} img_loop=$ISO_BASENAME archisobasedir=arch earlymodules=loop
  initrd (loop)/arch/boot/x86_64/archiso.img
}

EOT
}

function setup_freedos {
    install_memdisk
    copy_iso_data

    cat <<EOT > "$GRUB_PART_CFG"
menuentry '$ISO_BASENAME' {
  if [ \${grub_platform} = pc ]; then
    linux16 /memdisk raw
    initrd16 /$ISO_BASENAME
  else
    echo "FreeDOS only works with BIOS boot."
    sleep 3
  fi
}

EOT
}

function copy_iso_data {
    rsync --size-only --progress "$ISO_FILENAME" "${MULTIBOOT_MNT}/"
}

function install_memdisk {
    local dest="${MULTIBOOT_MNT}/memdisk"

    if [ -f "$dest" ]; then
        return
    fi

    for maybe in /usr/lib/syslinux/bios/memdisk /usr/lib/syslinux/memdisk; do
        if [ -f "$maybe" ]; then
            cp "$maybe" "$dest"
            return
        fi
    done

    curl --fail "$SYSLINUX_URL" | \
        tar xz --strip-components=3 -C "$MULTIBOOT_MNT" \
            "syslinux-${SYSLINUX_VERSION}/bios/memdisk/memdisk"
}

function delete_stray_grub_parts {
    pushd "${MULTIBOOT_MNT}/grub/parts" &> /dev/null
    comm -13 \
         <(find "${MULTIBOOT_MNT}/" -maxdepth 1 -type f -printf '%f\n'|sort) \
         <(find "${MULTIBOOT_MNT}/grub/parts/" -type f -printf '%f\n'|sort) \
         | xargs -d '\n' -r rm
    popd &> /dev/null
}

function print_grub_cfg {
    cat <<EOT
insmod all_video
insmod part_msdos
insmod progress
search --set=root --label $PARTITION_LABEL

EOT

    cat "${MULTIBOOT_MNT}/grub/parts/"*
}

CMD="cmd_${1:-}"
shift || true

if [ "$(type -t -- "$CMD")" = "function" ]; then
    "${CMD}" "$@"
else
    echo "Usage: $0 [format|mount|bootloader|add_iso|umount]"
    exit 1
fi