#!/bin/bash
USE_EFI=
die()
{
    printf 'Error: %s\n' "$*" >&2
    exit 1
}
KiB()
{
    printf %s $(("$1" * 1024))
}
MiB()
{
    KiB $(KiB "$1")
}

get_size()
{
    [ "$1" ] || return
    if [ -b "$1" ]
    then
        blockdev --getsize64 "$1"
    else
        stat -c'%s' "$1"
    fi
}

get_rootfs_size()
{
    rootfs_size=$(get_size "$ROOTFS")
}

initialize_target()
{
    if [ -b "$target" ]
    then return
    fi

    rm -f "$target"
    if [ "$ROOTFS" ]
    then
        get_rootfs_size
        fallocate -l $(($(KiB 17) + $(MiB $end_keys) + rootfs_size + $(KiB 640))) "$target"
    else
        fallocate -l $(($(KiB 17) + $(MiB $end_keys))) "$target"
    fi
}

# TODO: Use GUID type tags for each Samizdat partition type.
partition_target()
{
    parted "$target" -sm -a optimal                             \
           unit MiB                                             \
           mklabel gpt                                          \
           mkpart samizdat-efi  1 8                             \
           set 1 esp on                                         \
           mkpart samizdat-grub 8 64                            \
           set 2 bios_grub on                                   \
           mkpart samizdat-keys btrfs ${start_keys} ${end_keys}
    if [ -b "$target" ]
    then
        partx -u "$target"
    fi

    if [ "$ROOTFS" ]
    then
        parted "$target" -sm -a optimal                         \
               mkpart samizdat-rootfs btrfs ${end_keys}MiB 100%
    fi
}

sex()
{
    (
        set -x
        "$@"
    )
}

make_target_bootable()
{
    mkfs.btrfs -f --mixed "$part"3
    mkdir -p "$mnt"
    mount "$part"3 "$mnt"

    mkdir -p "$mnt"/boot/grub
    cp -avL "$GRUB_CONFIG" "$mnt"/boot/grub

    mkdir -p "$mnt"/boot/efi
    mkfs.vfat "$part"1
    if [ "$USE_EFI" ]
    then
        efi_dir="$mnt".efi
        mkdir -p "$efi_dir"
        mount "$part"1 "$mnt".efi
        mkdir "$efi_dir"/boot
        boot_dir=$efi_dir/boot
    else
        boot_dir=$mnt/boot
    fi

    # https://wiki.archlinux.org/index.php/Multiboot_USB_drive#Hybrid_UEFI_GPT_+_BIOS_GPT/MBR_boot
    # grub-install --target=x86_64-efi --removable --boot-directory=/mnt/boot --efi-directory=/mnt

    GRUB_LOGFILE=$(mktemp) || exit
    set -- eatmydata -- \
        grub-install ${VERBOSE:+ -v} \
        ${USE_EFI:+ --target=i386-efi} \
        ${USE_EFI:+ --recheck} \
        ${USE_EFI:+ --removable} \
        ${USE_EFI:+ --efi-directory "$efi_dir"} \
        --boot-directory "$boot_dir" \
        "$dev"
    echo "running: $* >$GRUB_LOGFILE 2>&1" >&2
    if "$@" >"$GRUB_LOGFILE" 2>&1
    then
        fail=
    else
        fail=y
    fi
    if [ "$USE_EFI" ]
    then
        umount -l "$mnt"/boot/efi || :
    fi
    umount -l "$mnt" || :
    if [ "$fail" ]
    then
        set -x
        tail -n5 "$GRUB_LOGFILE"
        exit 1
    else
        rm "$GRUB_LOGFILE"
    fi
    sync
}

copy_rootfs()
{
    if [ -z "$ROOTFS" ]
    then
        return
    fi
    [ -e "$ROOTFS" -a -b "$part"4 ]
    get_rootfs_size
    part_size=$(blockdev --getsize64 "$part"4)
    [ "$rootfs_size" -le "$part_size" ] || die "rootfs_size=$rootfs_size > part_size=$part_size"

    (
        set -x
        dd if="$ROOTFS" of="$part"4
        # dd if=rootfs/samizdat.patch.seed.btrfs of="$part"5
    )
}

acquire_target_block_devices()
{
    if [ -b "$target" ]
    then
        dev=$target
        part=$target
        mnt=/mnt/${target#/dev/}
        partx -u "$target"
    else
        trap release_target_block_devices EXIT
        losetup -f "$target"
        dev=$(losetup -j "$target" -O NAME --noheadings)

        btrfs device scan -u || :
        kpartx -vasas "$dev"
        btrfs device scan
        part=/dev/mapper/${dev##*/}p

        mnt=$target.mnt
    fi
}

release_target_block_devices()
{
    (
        set +e
        [ -z "$mnt" ] || ! mountpoint -q "$mnt" || umount "$mnt"
        case "$dev" in
            /dev/loop*)
                kpartx -d "$dev"
                losetup -d "$dev"
                ;;
        esac
        true
    )
    unset mnt dev
}

add_keys()
{
    rsync  -a  --info=STATS  "$GPG_INPUT_DIR"/ "$mnt"/gnupghome/
}

add_initrd()
{
    initrd_suffix=.samizdat
    mkdir -p "$mnt"/linux
    rsync  -aL --info=STATS "${1}vmlinuz${2}" "$mnt"/linux/vmlinuz
    rsync  -aL --info=STATS "${1}initrd.img${2}${2:+$initrd_suffix}" "$mnt"/linux/initrd.img
}

add_grub_cfg()
{
    mkdir -p "$mnt"/boot/grub
    cp -aL "$GRUB_CONFIG" "$mnt"/boot/grub
}

individualize()
{
    mkdir -p "$mnt"
    (
        set +e
        losetup -a >&2
        lsblk >&2
        mount|grep $PWD >&2
        : OK
    )
    mountpoint "$mnt" || mount "$part"3 "$mnt"
    add_keys
    add_initrd "$samizdat_linux_dir"/ "${version_suffix}"
    add_grub_cfg
}

globalize()
{
    initialize_target
    partition_target
    acquire_target_block_devices
    make_target_bootable
    copy_rootfs
}

sanity_checks()
{
    [[ $UID = 0 ]] || die "You are not root."
    for d in "$samizdat_linux_dir" "$GPG_INPUT_DIR"
    do
        [ -d "$d" ] || die "Not a directory: $d"
    done
}

find_mac()
{
    start_mac=$1
    for mac in $(ip link show | grep link/ether | (read _ mac _; echo $mac | tr : -)); do
        if [ "${mac%??}" = "${start_mac%??}" ]; then
            prefix=${mac%??}
            suffix=$(printf %x $(( 0x${mac##*-} + 1 )))
            MAC=${prefix}${suffix}
            return
        fi
    done
    MAC=$start_mac
}

boot_vm()
{
    boot_drive=$1
    installer_target=samizdat.disk.img

    if [ ! -e "$installer_target" ]
    then
        fallocate -l 10G "$installer_target"
    fi

    find_mac 52-54-00-12-34-56
    qemu-system-x86_64 \
        -enable-kvm \
        ${USE_EFI:+ -bios /usr/share/OVMF/OVMF_CODE.fd} \
        -smp 2 \
        -m 512 \
        -k en-us \
        -net nic,model=virtio,macaddr="$MAC" \
        -net tap,ifname=tap0,script=no,downscript=no \
        -curses \
        -drive file="$boot_drive",format=raw \
        -drive file="$installer_target",format=raw
}

do_build()
{
    if [ -b "$target" ]
    then
        globalize
        individualize
        release_target_block_devices
    elif [ -f "$target" -o ! -e "$target" ]
    then
         template=${target%.img}.template.img
         final=$target
         (
             target=$target.tmp
             if [ ! -e "$template" ]
             then
                 target="$template".tmp globalize
                 release_target_block_devices
                 mv -T "$template".tmp "$template"
             fi

             cp -T --reflink=always "$template" "$target"
             acquire_target_block_devices
             individualize
             release_target_block_devices

             mv -T "$target" "$final"
         )
    else
        die "Target must be a file or block device."
    fi
}

. samizdat-paths.sh || die 'samizdat-paths.sh not found'
samizdat_linux_dir=/

DEFAULT_ROOTFS=rootfs/samizdat.seed.btrfs
#DEFAULT_PATCHFS=rootfs/samizdat.patch.btrfs

if [ "$ROOTFS" = none ]
then
  ROOTFS=
elif [ "$ROOTFS" ]
then
  [ -f "$ROOTFS" ] || die "Not a file: $ROOTFS"
elif [ -f "$DEFAULT_ROOTFS" ]
then
  ROOTFS=$DEFAULT_ROOTFS
fi

if [ ! "$GPG_INPUT_DIR" ]
then
  for d in $GPG_INPUT_DIR /root/.gnupg /cdrom/gnupghome
  do
     [ -d "$d" ] || continue
     GPG_INPUT_DIR=$d
     break
  done
fi

: ${GRUB_CONFIG:=conf/grub.cfg}
start_keys=64 # megs
end_keys=256 # megs

if [ "$KERNEL_VERSION" ]
then
    version_suffix=-$KERNEL_VERSION
fi

sanity_checks
set -e

if [ "$1" = -n ]
then
    shift
    DO_BUILD=
else
    DO_BUILD=y
fi

default_target=boot-disk.img

USB=$(which usb || echo ./src/usb)

if [ "$1" ]
then
    if [ "$1" = USB ]
    then
        [ "$DO_BUILD" ] && arg= || arg=-n
        target=$($USB $arg) || exit 1
    elif [ -f "$1" ]
    then
        case "$1" in
            *.img) ;;
            *) die "Invalid image filename.  Use *.img";;
        esac
        target=$1
    fi
else
    target=$default_target
fi

if [ "$DO_BUILD" ]
then
    do_build
fi
boot_vm "$target"