#!/bin/bash

msg() { printf '%s: %s: %s\n' "$0" "$1" "$2" >&2; }
die() { msg Error "${*:-Exiting on fatal error.}"; exit 1; }
warn() { msg Warning "${*:-Something is wrong.}"; }
notice() { msg Notice "$*"; }
quietly() { "$@" >/dev/null 2>&1 || true; }
loudly() { (set -x; "$@"); }

validate_name()
{
        case "$1" in
        *[^a-zA-Z0-9_]*) false ;;
        *) true ;;
        esac
}

read_config_file()
{
        validate_name "$img" || { warn "invalid name: $img"; return 1; }
        while read
        do
                REPLY=${REPLY%%#*} # ignore comments
                [ "$REPLY" ] || continue # ignore empty lines
                k=${REPLY%%=*}
                v=${REPLY#*=}
                [ "$k" -a "$k" != "$REPLY" ] || return
                eval "conf_${img}_$k=\$v"
        done < "$img".conf
}

inquire_var() { _inquire_var "$img" "$1"; }
_inquire_var()
{
        local v
        v=conf_${1}_${2}
        v=${!v}
        if [ "$v" ]
        then
                eval "$2=\$v"
        else
                false
        fi
}

require_var() { _require_var "$img" "$1"; }
_require_var()
{
        _inquire_var "$@" || die "Missing required field '$2' for image file '$1'"
}

get_root_hash()
{
        sed -ne 's/^Root hash:[ \t]*//p' "$1"
}

require_exists()
{
        local f
        for f
        do
                [ -f "$f" ] || die "Not a file: $f"
        done
}

build_partition_image()
{
        if inquire_var rebuild
        then
                case "$rebuild" in
                        always|never|default) ;;
                        *) die "invalid value for field 'rebuild': $rebuild" ;;
                esac
        fi

        if [ "$rebuild" = 'always' ] || [ ! -e "$imgfile" -a "$rebuild" != 'never' ]
        then

                if [ -e "$imgfile" ]
                then
                        notice "Image file exists: $imgfile"
                fi

                case "$type" in
                        dm-verity-hashes|dm-verity-data)
                                require_var data_path
                                data_verity_path=../_build/${data_path##*/}.verity
                                require_exists "$data_path" "$data_verity_path" "$data_verity_path".log
                                root_hash=$(get_root_hash "$data_verity_path".log)
                                [ ${#root_hash} = 64 ]
                        ;;
                        *)
                                require_var allocation
                                tmp=$imgfile~tmp
                                fallocate -l "$allocation" "$tmp"
                        ;;
                esac

                case "$type" in
                        bios-grub)            ;;
                        efi-system-partition) mkfs.fat -F 32 -I "$tmp" || die "mkfs.vfat failed" ;;
                        boot|samizdat-keys)   mkfs.btrfs -q "$tmp" || die "mkfs.btrfs failed" ;;
                        dm-verity-data)
                                partuuid=${root_hash:0:32}
                                set_var partuuid
                                cp -f -T --reflink "$data_path" "$builddir"/"$partuuid"
                                ln -sfT "$partuuid" "$tmp"
                        ;;
                        dm-verity-hashes)
                                partuuid=${root_hash:32:32}
                                set_var partuuid
                                cp -f -T --reflink "$data_verity_path" "$builddir"/"$partuuid"
                                ln -sfT "$partuuid" "$tmp"
                        ;;
                esac
                mv -T "$tmp" "$imgfile"
                notice "Successfully wrote $imgfile"
        fi
}

iterate_partitions()
{
        local f
        for f in part*.conf
        do
                notice "Processing $f"
                img=${f%.conf}

                read_config_file || warn "Received error return from command: read_config_file $img"
                require_var name

                require_var type
                case "$type" in
                        dm-verity-hashes|dm-verity-data)
                            require_var data_path
                            [ ! "$SKIP_ROOTFS_COPY" ] || continue
                            ;;
                        efi-system-partition|bios-grub|boot|samizdat-*|partition-table) ;;
                        *) die "invalid type: $type" ;;
                esac

                imgfile=$builddir/$img
                "$@" || return
        done
}

set_var()
{
        eval "conf_${img}_${1}=\$${1}"
}

format_guid()
{
    set -- "$(tr -dc 0-9a-fA-F <<< "$1")"
    [ ${#1} = 32 ]
    printf '%s\n' "${1:0:8}-${1:8:4}-${1:12:4}-${1:16:4}-${1:20:12}"
}

get_partition_size()
{
    case "$type" in
        dm-verity-hashes|dm-verity-data)
            stat -L -c '%s' "$builddir"/"${f%.conf}"
            ;;
        *)
            require_var allocation
            case "$allocation" in
                *K) echo $(( ${allocation%?} * 1024 )) ;;
                *M) echo $(( ${allocation%?} * 1024 * 1024 )) ;;
                *G) echo $(( ${allocation%?} * 1024 * 1024 * 1024 )) ;;
                *) return 1 ;;
            esac
            ;;
    esac
}

create_ptable_conf()
{
        unset partuuid
        inquire_var partuuid
        devsz=$(get_partition_size) || return

        [ "$start" -a "$devsz" ] || return

        [ "$((devsz % 512))" -eq 0 ]

        start_sectors=$((start / 512))
        size_sectors=$((devsz / 512))

        typecode=0FC63DAF-8483-4772-8E79-3D69D8477DE4
        case "$type" in
                partition-table) start=$((start + devsz)); return;;
                efi-system-partition) typecode=C12A7328-F81F-11D2-BA4B-00A0C93EC93B ;;
                dm-verity-data|dm-verity-hashes)
                    case "$name" in
                        samizdat-rootfs|samizdat-root-patch) typecode=4f68bce3-e8cd-4db1-96e7-fbcaf984b709 ;;
                        samizdat-root-patch-verity) typecode=2c7357ed-ebd2-46d9-aec1-23d437ec2bf5 ;;
                        samizdat-root-seed) ;; # keep default
                        samizdat-root-seed-verity) ;; # keep default
                    esac
                    ;;
                root) typecode=4f68bce3-e8cd-4db1-96e7-fbcaf984b709 ;;
                bios-grub)
                    typecode=21686148-6449-6E6F-744E-656564454649
                    printf 'start=1, size=%d, type=ee\n' \
                           "$((start_sectors - 1))" \
                           >> "$DOS_TABLE_FILE"
                    printf 'start=%d, size=%d, type=83, bootable\n' \
                           "$start_sectors" \
                           "$size_sectors" \
                           >> "$DOS_TABLE_FILE"
                    printf 'start=%d, type=ee\n' \
                           "$((start_sectors + size_sectors))" \
                           >> "$DOS_TABLE_FILE"
                    ;;
        esac

        (
            exec >> "$GPT_TABLE_FILE"
            printf '%d: start=%d, size=%d, type=%s, name="%s"' \
                    "$i" \
                    "$start_sectors" \
                    "$size_sectors" \
                    "$typecode" \
                    "$name"
            if [ "$partuuid" ]
            then
                printf ', uuid=%s' "$(format_guid "$partuuid")"
            fi
            echo
        )
        let ++i
        start=$((start + devsz))
}

pee_on_table()
{
        GPT_TABLE_FILE=$builddir/table.gpt
        DOS_TABLE_FILE=$builddir/table.mbr

        printf 'label: gpt\n\n' > "$GPT_TABLE_FILE"
        : > "$DOS_TABLE_FILE"

        local dev="$1"
        i=1
        start=0
        iterate_partitions create_ptable_conf

        loudly $sudo sfdisk --no-tell-kernel "$dev" < "$GPT_TABLE_FILE" || return
        loudly $sudo sfdisk --no-tell-kernel -Y dos "$dev" < "$DOS_TABLE_FILE" || return
}

clone_parts_to_target()
{

        f=$(readlink -e "$builddir"/"${f%.conf}") || return
        samizdat-ficlonerange "$f" "$target" 0 0 0
}

cleanup()
{
    for f in part*.conf
    do
        mnt=${f%.conf}.mnt
        if mountpoint -q "$mnt"
        then
            loudly $sudo umount "$mnt"
            $sudo rmdir "$mnt"
        fi
    done
    if [ "$whole" ]
    then
        loudly $sudo kpartx -sd "$whole"
        loudly $sudo losetup -d "$whole"
    fi
}

mount_part()
{
    local source_dev="$1"
    case "$type" in
        efi-system-partition|boot|samizdat-keys) ;;
        dm-verity-hashes|partition-table|bios-grub) return ;;
        *) notice "Not mounting $name"; return ;;
    esac

    img=${f%.conf}
    mnt=${f%.conf}.mnt

    mkdir -p "$mnt"
    loudly $sudo mount "$source_dev" "$mnt"
}

copy_data_to_mounted_target_filesystems()
{
    mount_part "/dev/mapper/${whole#/dev/}p${img#part}"
    case "$type" in
        boot)
            BOOT_DIR=$mnt
            $sudo rsync -a --info=STATS /boot/ "$mnt"/
            $sudo systemd-run -p BindPaths="$(realpath -e "$mnt"):/boot" --wait update-grub
            ;;
        samizdat-keys)
            ;;
        efi-system-partition)
            EFI_DIR=$mnt
            $sudo mkdir "$EFI_DIR"/EFI "$EFI_DIR"/EFI/BOOT
            ;;
    esac
}

sanity_check()
{
        grep -q samizdat "$BOOT_DIR"/grub/grub.cfg
}

shopt -s nullglob
PATH=/sbin:$PATH

builddir=_build

if [ "$UID" = 0 ]
then sudo=
else sudo=sudo
fi

SKIP_ROOTFS_COPY=
if [ "$1" = 'key' ]
then
    SKIP_ROOTFS_COPY=y
fi

set -e

mkdir -p "$builddir"

if [ "$SKIP_ROOTFS_COPY" ]
then
    target=$builddir/key.img
else
    target=$builddir/whole.img
fi






if whole=$(losetup -j "$target" -O NAME --noheadings)
then
    cleanup
    whole=
fi
truncate -s0 "$target"
iterate_partitions build_partition_image




iterate_partitions clone_parts_to_target
$sudo losetup -L -f "$target"
whole=$(losetup -j "$target" -O NAME --noheadings)







pee_on_table "$whole"





trap cleanup EXIT
$sudo kpartx -su "$whole"

iterate_partitions copy_data_to_mounted_target_filesystems
sanity_check
time loudly $sudo eatmydata -- grub-install --boot-directory="$BOOT_DIR" "$whole" --target=i386-pc
time loudly $sudo eatmydata -- grub-install --boot-directory="$BOOT_DIR" "$whole" --target=x86_64-efi --removable --efi-directory="$EFI_DIR" || true

valdik_efi_dir=../import-grub-bootx64-efi/3
$sudo cp -v --preserve=timestamps -t "$EFI_DIR"/EFI/BOOT/ -- "$valdik_efi_dir"/BOOTX64.EFI