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

get_rootfs_size()
{
    [ "$ROOTFS" ] || return
    if [ -b "$ROOTFS" ]
    then
        rootfs_size=$(blockdev --getsize64 "$ROOTFS")
    else
        rootfs_size=$(stat -c%s "$ROOTFS")
    fi
}

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-grub 1 8                             \
           set 1 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
}

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

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

    eatmydata -- grub-install -v --boot-directory "$mnt"/boot "$dev"
    sync
}

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

    dd if="$ROOTFS" of="$part"3
}

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)

        kpartx -vasas "$dev"
        part=/dev/mapper/${dev##*/}p

        mnt=$target.mnt
    fi
}

release_target_block_devices()
{
    (
        set +e
        [ -z "$mnt" ] || ! mountpoint "$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"
    mountpoint "$mnt" || mount "$part"2 "$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 \
        -smp 2 \
        -m 512 \
        -k en-us \
        -net nic,model=virtio,macaddr="$MAC" \
        -net tap,ifname=tap0,script=no,downscript=no \
        -vga qxl \
        -drive file="$boot_drive",format=raw \
        -drive file="$installer_target",format=raw
}

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



#ROOTFS=rootfs/samizdat.seed.btrfs
#PATCHFS=rootfs/samizdat.patch.btrfs
#GPG_INPUT_DIR=/root/.gnupg

: ${GPG_INPUT_DIR:=/cdrom/gnupghome}
: ${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" = boot ]
then
    boot_vm boot-disk.img
elif [ "$1" ]
then
    [ -b "$1" ] || die "Not a block device: $1"
    target=$1
    globalize
    individualize
    release_target_block_devices
else
    template=boot-disk.template.img
    if [ ! -e "$template" ]
    then
        target="$template".tmp globalize
        release_target_block_devices
        mv -T "$template".tmp "$template"
    fi

    target=boot-disk.img.tmp
    cp -T --reflink=always "$template" "$target"
    acquire_target_block_devices
    individualize
    release_target_block_devices

    final=${target%.tmp}
    mv -T "$target" "$final"
    target=$final
fi
boot_vm "$target"