#!/bin/sh

# TODO: Make sure fdisk from busybox is NOT used, it's counting sectors incorrectly (?!)
# TODO: Check connectivity with API server

BUTTERKNIFE_POOL_MOUNTPOINT=/var/lib/butterknife/pool
TARGET_MOUNTPOINT=/mnt/target

AGENT="Butterknife-Provisioning-Image/0.2"
URL_LOCAL=http://butterknife


#######################################
### Check for presence of utilities ###
#######################################
for util in btrfs jq fdisk find ntfsresize udp-sender udp-receiver ntpdate curl ms-sys mktemp test true sort uniq dig; do
    if [ -z "$(which $util)" ]; then
        echo "Butterknife was unable to locate $util," \
            "are you sure the provisioning image was compiled properly?"
        exit 253
    else
        echo "Found $util..."
    fi
done

#for i in $(cat /proc/cmdline); do
#    case i in
#        bk_debug)
            set -x # Echo
#        ;;
#    esac
#done


set -e # Bail on error

#######################################
### Transfer method selection phase ###
#######################################

if [ -z $BUTTERKNIFE_TRANSFER_METHOD ]; then
    BUTTERKNIFE_TRANSFER_METHOD=$(dialog --menu "Select transfer method" 0 0 0 \
        http "HTTP-only" \
        multicast "Multicast receive" \
        tee "Multicast via HTTP and write" \
        proxy "Only proxy HTTP to multicast" 2>&1 >$(tty))
    clear
fi


##################################
### Harddisk preparation phase ###
##################################

case $BUTTERKNIFE_TRANSFER_METHOD in
    http|multicast|tee)
        #############################
        ### Target disk selection ###
        #############################

        if [ -z $bk_disk_slug ]; then
            bk_disk_slug=$(butterknife-select-disk 2>&1 >$(tty))
        fi
        BUTTERKNIFE_DISK=/dev/$bk_disk_slug
        
        if [ -z $BUTTERKNIFE_PARTITIONING_METHOD ]; then

            BUTTERKNIFE_PARTITIONING_METHOD=$(dialog --menu "Partitioning $BUTTERKNIFE_DISK" 0 0 0 \
                receive         "Receive into existing btrfs filesystem" \
                unpartitioned   "Use unpartitioned area" \
                resize          "Resize last partition" \
                reformat        "Reformat partition" \
                mbr             "Wipe <2TB disk" \
                gpt             "Wipe 2TB+ disk" \
                efi             "Wipe and do EFI install" \
                2>&1 >$(tty))
            clear
        fi

        # TODO: deploy          "Deploy received template" \
        # TODO: postinstall     "Run postinstall scripts (reinstall GRUB)" \
        case $BUTTERKNIFE_PARTITIONING_METHOD in
            unpartitioned)
                echo "Attempting to create new partition in unpartitioned space"
                echo -e "n\np\n\n\n\nw" | fdisk $BUTTERKNIFE_DISK
            ;;
            resize)
                NTFS_PARTITION=$(ls $BUTTERKNIFE_DISK? | tail -n 1)
                # TODO: Assert last one is NTFS
                # TODO: Suggested size heuristics
                MINSIZE=$(ntfsresize $NTFS_PARTITION -m | grep Minsize | cut -d ':' -f 2)
                SUGGESTED=${MINSIZE}M
                SIZE=$(dialog --inputbox "Enter new filesystem size of at least ${MINSIZE}M" 0 0 $SUGGESTED 2>&1 >$(tty))
                clear
                ntfsresize -s $SIZE $NTFS_PARTITION
                echo -e "d\n\nw" | fdisk $BUTTERKNIFE_DISK # Remove last partition
                echo -e "n\np\n\n\n+$SIZE\nt\n\n7\nw" | fdisk $BUTTERKNIFE_DISK # Re-create NTFS
                echo -e "n\np\n\n\n\nw" | fdisk $BUTTERKNIFE_DISK # Create partition for btrfs
            ;;
            mbr)
                echo "Purging whole disk"
                echo -e "o\nn\np\n\n\n\na\n1\nw" | fdisk $BUTTERKNIFE_DISK
            ;;
            gpt)
                sgdisk $BUTTERKNIFE_DISK -o -g \
                    -n 1::+2MB -t 1:ef02 -c 1:"BIOS Boot Partition" \
                    -n 2       -t 2:8300 -c 2:"Butterknife pool" -p
            ;;
            efi)
                sgdisk $BUTTERKNIFE_DISK -o -g \
                    -n 1::+100MB -t 1:ef00 -c 1:"EFI System Partition" \
                    -n 2         -t 2:8300 -c 2:"Butterknife pool" -p
            ;;
            reformat|receive)
                # NOOP
            ;;
            *)
                echo "Invalid partitioning method $BUTTERKNIFE_PARTITIONING_METHOD"
                exit 255
            ;;
        esac


        ############################################
        ### Target partition determination phase ###
        ############################################
        case $BUTTERKNIFE_PARTITIONING_METHOD in
            reformat|receive)
                # Dialog to select partition for reformat or receive
                for partition in $BUTTERKNIFE_DISK?; do
                    partition_slug=$(echo $BUTTERKNIFE_PARTITION | cut -d "/" -f 3)
                    sector_count=$(cat /sys/block/$bk_disk_slug/$BUTTERKNIFE_PARTITION_slug/size)
                    sector_size=$(cat /sys/block/$bk_disk_slug/queue/hw_sector_size)
                    size=$(expr $sector_count / 1000000 \* $sector_size / 1000 || true)G
                    if [ $size == "0G" ]; then
                        size=$(expr $sector_count / 1000 \* $sector_size / 1000 || true)M
                    fi
                    echo "$BUTTERKNIFE_PARTITION \"$size\"";
                done > /tmp/partitions

                BUTTERKNIFE_PARTITION=$(dialog \
                    --menu "Target partition" 0 0 0 \
                    --file /tmp/partitions \
                    2>&1 >$(tty))
                clear
            ;;
            unpartitioned|resize|mbr|gpt|efi)
                # Assume last partition
                BUTTERKNIFE_PARTITION=$(ls $BUTTERKNIFE_DISK? | tail -n 1)
            ;;
            *)
                echo "Invalid partitioning method $BUTTERKNIFE_PARTITIONING_METHOD"
                exit 255
            ;;
        esac

        PARTITION_SLUG=$(echo $BUTTERKNIFE_PARTITION | cut -d "/" -f 3)


        ########################################
        ### Target filesystem creation phase ###
        ########################################
        case $BUTTERKNIFE_PARTITIONING_METHOD in
            efi)
                mkfs.vfat ${BUTTERKNIFE_DISK}1
            ;;
        esac

        case $BUTTERKNIFE_PARTITIONING_METHOD in
            mbr|gpt|efi|reformat|unpartitioned|resize)
                echo "Creating clean btrfs filesystem on $PARTITON"
                mkfs.btrfs -f $BUTTERKNIFE_PARTITION
            ;;
        esac

        # Attempt to mount target directory
        mkdir -p $BUTTERKNIFE_POOL_MOUNTPOINT
        mount $BUTTERKNIFE_PARTITION $BUTTERKNIFE_POOL_MOUNTPOINT -o subvol=/ -t btrfs
        if [ $? -ne 0 ]; then
            dialog --msgbox "Mounting $BUTTERKNIFE_PARTITION at $BUTTERKNIFE_POOL_MOUNTPOINT failed, are you sure kernel has btrfs support built-in?" 0 0
            exit 255
        fi

        ################
        ### Clean up ###
        ################

        for subvol in $(ls $BUTTERKNIFE_POOL_MOUNTPOINT |  (grep "^@template:" || true)); do
            set +e
            touch $BUTTERKNIFE_POOL_MOUNTPOINT/$subvol/.test
            if [ $? -eq 0 ]; then
                set -e
                btrfs subvol delete $BUTTERKNIFE_POOL_MOUNTPOINT/$subvol
            fi
            set -e
        done
    ;;
    *)
        BUTTERKNIFE_PARTITIONING_METHOD="pass"
    ;;
esac

##############################
### Determine architecture ###
##############################

bk_arch=$(uname -m | sed 's/^i.86$/x86/')

case $BUTTERKNIFE_TRANSFER_METHOD in
    http|tee|proxy)
        ##############################
        ### Server selection phase ###
        ##############################
        if [ -z $bk_url ]; then

            bk_url=$(dialog --menu "Select server" 0 0 0 \
                mdns:// "Autodiscover" \
                $URL_LOCAL                       "Manually enter" \
                https://butterknife.k-space.ee   "K-SPACE MTÜ" 2>&1 >$(tty))

            if [ "$bk_url" == "mdns://" ]; then
                butterknife-discover > /tmp/discovered_servers
                bk_url=$(dialog --menu "Select one of discovered servers" \
                    0 0 0 --file /tmp/discovered_servers 2>&1 >$(tty))
            elif [ "$bk_url" == $URL_LOCAL ]; then
                bk_url=$(dialog --inputbox "Manually enter the URL of Butterknife server" \
                    0 0 $URL_LOCAL 2>&1 >$(tty))
            fi
            clear
        fi


        ################################
        ### Template selection phase ###
        ################################

        if [ -z $bk_template ]; then
            # Fetch template list
            curl -A $AGENT -s $bk_url/api/template \
                | jq '.templates[] | .namespace + "." + .identifier + " \"" + .description + "\""' -r \
                > /tmp/available_templates

            bk_template=$(dialog \
                --menu "Select template to deploy" 0 0 0 \
                --file /tmp/available_templates \
                2>&1 >$(tty))
            clear

        fi

        ###############################
        ### Version selection phase ###
        ###############################
        if [ -z $bk_version ]; then
            # Fetch version list
            curl -A $AGENT -s $bk_url/api/template/$bk_template/arch/$bk_arch/version \
                > /tmp/available_versions.json

            cat /tmp/available_versions.json \
                | jq '.versions[] | .identifier + " \"" + .comment + "\""' -r \
                | head -n 100 \
                > /tmp/available_versions

            bk_version=$(dialog \
                --menu "Select version to deploy" 0 0 0 \
                --file /tmp/available_versions \
                2>&1 >$(tty))
            clear
        fi

        BUTTERKNIFE_TEMPLATE_SUBVOL="@template:$bk_template:$bk_arch:$bk_version"



        #####################################
        ### Stream URL construction phase ###
        #####################################

        # Build btrfs-stream URL
        STREAM="$bk_url/$BUTTERKNIFE_TEMPLATE_SUBVOL"
    ;;
esac


##############################################
### Allow differential versions using HTTP ###
##############################################

case $BUTTERKNIFE_TRANSFER_METHOD in
    http)
        # Determine differential version parent
        cat /tmp/available_versions.json | jq -r '.versions[] .identifier' > /tmp/available_version_names
        ls $BUTTERKNIFE_POOL_MOUNTPOINT | (grep "^@template:$bk_template:$bk_arch:" || true) | cut -d ":" -f 4 > /tmp/local_version_names
        PARENT=$(cat /tmp/local_version_names /tmp/available_version_names | sort | grep -v $bk_version | uniq -d | sort -t p -k 2n | tail -n 1)
        if [ -z $PARENT ]; then
            echo "Could not determine parent, falling back to full snapshot"
        else
            STREAM="$STREAM?parent=@template:$bk_template:$bk_arch:$PARENT"
        fi
        echo "Final URL is $STREAM"
    ;;
esac

####################################################
### Enable compression if we're going over HTTPS ###
####################################################

case $STREAM in
    https://*)
        STREAM="--compressed $STREAM"
    ;;
esac


######################
### Transfer phase ###
######################

case $BUTTERKNIFE_TRANSFER_METHOD in
    multicast)
        ls $BUTTERKNIFE_POOL_MOUNTPOINT | (grep "^@template:" || true) > /tmp/local_templates
        udp-receiver --nokbd | btrfs receive $BUTTERKNIFE_POOL_MOUNTPOINT
        # Heuristics to determine name of received snapshot
        ls $BUTTERKNIFE_POOL_MOUNTPOINT | (grep "^@template:" || true) > /tmp/new_templates
        BUTTERKNIFE_TEMPLATE_SUBVOL=$(cat /tmp/local_templates /tmp/new_templates | sort | uniq -u)
        # TODO: Break here if we got garbage
        bk_template=$(echo $BUTTERKNIFE_TEMPLATE_SUBVOL | cut -d ":" -f 2)
        bk_arch=$(echo $BUTTERKNIFE_TEMPLATE_SUBVOL | cut -d ":" -f 3)
        bk_version=$(echo $BUTTERKNIFE_TEMPLATE_SUBVOL | cut -d ":" -f 4)
    ;;
    http)
        curl -A $AGENT $STREAM | btrfs receive $BUTTERKNIFE_POOL_MOUNTPOINT
    ;;
    tee)
        dialog --msgbox "Press enter once all the other machines are ready to receive" 0 0
        mkfifo /tmp/multicast_stream /tmp/local_stream
        cat /tmp/local_stream | btrfs receive $BUTTERKNIFE_POOL_MOUNTPOINT &
        sleep 1
        udp-sender --nokbd --no-progress --min-receivers 1 --min-wait 5 /tmp/multicast_stream &
        sleep 1
        curl -A $AGENT  -s $STREAM | tee /tmp/multicast_stream > /tmp/local_stream
        sleep 2
        # TODO: Ensure btrfs receive has finished
    ;;
    proxy)
        dialog --msgbox "Press enter once all the other machines are ready to receive" 0 0
        curl -A $AGENT -s $STREAM \
            | udp-sender --nokbd --min-receivers 1 --min-wait 5
    ;;
esac

sync

case $BUTTERKNIFE_PARTITIONING_METHOD in
    pass)
        echo "Skipping template deployment"
    ;;
    *)
        #################################
        ### Template deployment phase ###
        #################################

        BUTTERKNIFE_DEPLOY_SUBVOL="@root:$bk_template:$bk_arch:$bk_version"
        btrfs subvolume snapshot \
            $BUTTERKNIFE_POOL_MOUNTPOINT/$BUTTERKNIFE_TEMPLATE_SUBVOL \
            $BUTTERKNIFE_POOL_MOUNTPOINT/$BUTTERKNIFE_DEPLOY_SUBVOL

        # Symlink @root:active to current deployment subvol
        rm -f $BUTTERKNIFE_POOL_MOUNTPOINT/@root:active
        ln -s \
            $BUTTERKNIFE_DEPLOY_SUBVOL \
            $BUTTERKNIFE_POOL_MOUNTPOINT/@root:active


        ###############################
        ### Run post-deploy scripts ###
        ###############################

        # Mount deployment subvolume at target directory
        mkdir -p $TARGET_MOUNTPOINT
        mount -o subvol=$BUTTERKNIFE_DEPLOY_SUBVOL \
            $BUTTERKNIFE_PARTITION \
            $TARGET_MOUNTPOINT

        # Mount pool also for chroot
        mkdir -p $TARGET_MOUNTPOINT$BUTTERKNIFE_POOL_MOUNTPOINT
        mount -o subvol=/ \
            $BUTTERKNIFE_PARTITION \
            $TARGET_MOUNTPOINT$BUTTERKNIFE_POOL_MOUNTPOINT

        # Mount stuff for chroot
        mount --bind /dev/ $TARGET_MOUNTPOINT/dev/
        mount --bind /sys/ $TARGET_MOUNTPOINT/sys/
        mount --bind /proc/ $TARGET_MOUNTPOINT/proc/
        mount --bind /run/ $TARGET_MOUNTPOINT/run/
        mount none $TARGET_MOUNTPOINT/tmp/ -t tmpfs
        
        case $BUTTERKNIFE_PARTITIONING_METHOD in
            efi)
                mount ${BUTTERKNIFE_DISK}1 $TARGET_MOUNTPOINT/boot/efi
            ;;
        esac

        # Export variables for postinstall scripts
        export BUTTERKNIFE_DOMAIN
        export BUTTERKNIFE_TEMPLATE_SUBVOL
        export BUTTERKNIFE_DEPLOY_SUBVOL
        export BUTTERKNIFE_PARTITION
        export BUTTERKNIFE_DISK
        export BUTTERKNIFE_POOL_MOUNTPOINT
        export BUTTERKNIFE_PARTITIONING_METHOD
        export BUTTERKNIFE_TRANSFER_METHOD
        export BUTTERKNIFE_POOL_UUID=$(blkid -s UUID -o value $BUTTERKNIFE_PARTITION)

        # Copy DNS config
        mkdir -p /run/resolvconf
        mkdir -p /run/systemd/resolve
        cat /etc/resolv.conf > $TARGET_MOUNTPOINT/etc/resolv.conf

        # Export sensible PATH
        export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

        # Run postinstall scripts, presumably in sorted order
        chroot $TARGET_MOUNTPOINT butterknife-deploy
        chroot $TARGET_MOUNTPOINT butterknife-maintenance

        # Be forgiving from now on
        set +e

        # Unmount directories
        echo "Unmounting filesystems"
        umount -a
        echo "Flushing buffers"
        sync
        sleep 1
        echo "Rebooting machine"
        reboot -f
    ;;
esac