#!/bin/bash

set -eu

PARTITION_LABEL_ESP="multibt-esp"
PARTITION_LABEL_DATA="multiboot"
MULTIBOOT_MNT="${MULTIBOOT_MNT:-/mnt/multiboot}"

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

    sudo -k
    DISK_DEVICE="$1"
    print_sfdisk_command | sudo sfdisk --wipe always --wipe-partitions always "$DISK_DEVICE"
    udevadm settle
    sudo mkfs.vfat -n "${PARTITION_LABEL_ESP}" "/dev/disk/by-partlabel/${PARTITION_LABEL_ESP}"
    sudo mkfs.ext4 -L "${PARTITION_LABEL_DATA}" -E "root_owner=$(id -u):$(id -g)" "/dev/disk/by-partlabel/${PARTITION_LABEL_DATA}"
}

function print_sfdisk_command {
    cat <<EOT
label: gpt
unit: sectors

size=10M, type=uefi, name="$PARTITION_LABEL_ESP"
size=1M, type=21686148-6449-6E6F-744E-656564454649
type=linux, name="$PARTITION_LABEL_DATA"
EOT
}

function cmd_grub {
    DISK_DEVICE="$(findmnt -n -o source "$MULTIBOOT_MNT" | 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
    install_grub_cfg
}

function cmd_grubcfg {
    set -x
    install_grub_cfg
}

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

function install_grub_efi {
    for arch in i386-efi x86_64-efi; do
        sudo grub-install \
            --target="$arch" \
            --no-nvram \
            --removable \
            --efi-directory="${MULTIBOOT_MNT}/EFI" \
            --boot-directory="$MULTIBOOT_MNT" \
            "$DISK_DEVICE"
    done
}

function install_grub_cfg {
    if [[ -w "${MULTIBOOT_MNT}/grub/" ]]; then
        # We already have write access, no need to use sudo
        print_grub_cfg > "${MULTIBOOT_MNT}/grub/grub.cfg"
    else
        print_grub_cfg | sudo tee "${MULTIBOOT_MNT}/grub/grub.cfg" > /dev/null
    fi
}

function cmd_mount {
    set -x

    while sudo umount "/dev/disk/by-partlabel/${PARTITION_LABEL_ESP}" &> /dev/null; do true; done
    while sudo umount "/dev/disk/by-partlabel/${PARTITION_LABEL_DATA}" &> /dev/null; do true; done
    sudo mkdir -p "$MULTIBOOT_MNT"
    sudo mount "/dev/disk/by-partlabel/${PARTITION_LABEL_DATA}" "$MULTIBOOT_MNT"
    mkdir -p "${MULTIBOOT_MNT}/EFI"
    sudo mount "/dev/disk/by-partlabel/${PARTITION_LABEL_ESP}" "${MULTIBOOT_MNT}/EFI"
}

function cmd_umount {
    set -x

    sudo umount "${MULTIBOOT_MNT}/EFI" || true
    sudo umount "$MULTIBOOT_MNT" || true
    sudo rmdir "$MULTIBOOT_MNT"
}

function cmd_freedos {
    set -x

    local SYSLINUX_VERSION="6.03"
    local SYSLINUX_URL="https://www.kernel.org/pub/linux/utils/boot/syslinux/syslinux-${SYSLINUX_VERSION}.tar.gz"
    local FREEDOS_URL="https://www.ibiblio.org/pub/micro/pc-stuff/freedos/files/distributions/1.2/official/FD12LITE.zip"

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

    curl -fL "$FREEDOS_URL" > "${MULTIBOOT_MNT}/FD12LITE.zip"
}

function cmd_memtest {
    curl -fL -o "${MULTIBOOT_MNT}/memtest.tmp.zip" "https://memtest.org/download/v6.20/mt86plus_6.20_64.grub.iso.zip"
    unzip -d "$MULTIBOOT_MNT" "${MULTIBOOT_MNT}/memtest.tmp.zip"
    rm "${MULTIBOOT_MNT}/memtest.tmp.zip"
}

function print_grub_cfg {
    cat <<EOT
search --set=root --label $PARTITION_LABEL_DATA

insmod all_video
insmod gfxterm
loadfont unicode
set gfxmode=1024x768
terminal_output gfxterm

insmod part_msdos
insmod progress
insmod regexp

set maybe_quiet='quiet splash'
set maybe_to_ram=''

menuentry "! Copy ISO image to ram before booting" {
  # copytoram is used by arch
  # toram is used by casper based images (tails, Ubuntu, etc)
  set maybe_to_ram="copytoram toram"
}

menuentry "! Verbose" {
  set maybe_quiet=''
}

function setup_arch {
  menuentry "\$1" {
    loopback loop \$1
    echo "Loading kernel..."
    linux (loop)/arch/boot/x86_64/vmlinuz-* img_label=${PARTITION_LABEL_DATA} img_loop=\$1 archisobasedir=arch earlymodules=loop \$maybe_to_ram \$maybe_quiet
    echo "Loading initrd (and microcode if they exist)..."
    initrd (loop)/arch/boot/*.img (loop)/arch/boot/x86_64/initramfs-*.img
  }
}
for iso in /archlinux-*.iso /isos/archlinux-*.iso; do
  if [ -f "\$iso" ]; then
    setup_arch \$iso
  fi
done

function setup_debian {
  menuentry "\$1" {
    loopback loop \$1
    echo "Loading kernel..."
    linux (loop)/live/vmlinuz* boot=live components findiso=\$1 \$maybe_to_ram \$maybe_quiet
    echo "Loading initrd..."
    initrd (loop)/live/initrd*
  }
}
for iso in /debian-live-*.iso /isos/debian-live-*.iso; do
  if [ -f "\$iso" ]; then
    setup_debian \$iso
  fi
done

if [ -f /memdisk -a -f /FD12LITE.zip ]; then
  menuentry /FD12LITE.zip {
    if [ \${grub_platform} = pc ]; then
      linux16 /memdisk raw
      initrd16 /FD12LITE.zip
    else
      echo "FreeDOS only works with BIOS boot."
      sleep 3
    fi
  }
fi


function setup_memtest {
  menuentry "\$1" {
    loopback loop \$1
    if [ \${grub_platform} = pc ]; then
      linux (loop)/boot/memtest
    else
      linux (loop)/EFI/BOOT/memtest
    fi
  }
}
for iso in /mt86plus*.grub.iso /isos/mt86plus*.grub.iso; do
  if [ -f "\$iso" ]; then
    setup_memtest \$iso
  fi
done


function setup_fedora {
  menuentry "\$1" {
    loopback loop \$1
    probe -s iso_label -l (loop)
    echo "Loading kernel..."
    linux (loop)/images/pxeboot/vmlinuz root=live:CDLABEL=\$iso_label rd.live.image iso-scan/filename=\$1 \$maybe_quiet
    echo "Loading initrd..."
    initrd (loop)/images/pxeboot/initrd.img
  }
}
for iso in /Fedora-Workstation-Live-*.iso /isos/Fedora-Workstation-Live-*.iso; do
  if [ -f "\$iso" ]; then
    setup_fedora \$iso
  fi
done

function setup_ubuntu {
  menuentry "\$1" {
    loopback loop \$1
    set maybe_layerfs_path=''
    for f in minimal.standard.live.squashfs; do
      if [ -f "(loop)/casper/\$f" ]; then
        echo " \$f"
        set maybe_layerfs_path="layerfs-path=\$f"
        echo "Setting \$maybe_layerfs_path"
      fi
    done
    echo "Loading kernel..."
    linux (loop)/casper/vmlinuz* \$maybe_layerfs_path boot=casper iso-scan/filename=\$1 \$maybe_to_ram \$maybe_quiet
    echo "Loading initrd..."
    initrd (loop)/casper/initrd*
  }
}
for iso in /ubuntu-*-desktop-*.iso /isos/ubuntu-*-desktop-*.iso; do
  if [ -f "\$iso" ]; then
    setup_ubuntu \$iso
  fi
done

function setup_tails {
  menuentry "\$1" {
    loopback loop \$1
    echo "Loading kernel..."
    linux (loop)/live/vmlinuz* initrd=/live/initrd.img boot=live config iso-scan/filename=\$1 findiso=\$1 nopersistence noprompt timezone=Etc/UTC noautologin module=Tails slab_nomerge slub_debug=FZP mce=0 vsyscall=none page_poison=1 init_on_free=1 mds=full,nosmt \$maybe_to_ram \$maybe_quiet
    echo "Loading initrd..."
    initrd (loop)/live/initrd*
  }
}
for iso in /tails-*.iso /isos/tails-*.iso; do
  if [ -f "\$iso" ]; then
    setup_tails \$iso
  fi
done

EOT
}

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

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