#!/bin/bash

set -x
set -e -u

script_path="$(realpath "$0")"
script_dir="$(dirname "$script_path")"

old_ro_rootfs_mountpoint="/media/root-ro"
old_rw_overlay_mountpoint="/media/root-rw"

new_ro_rootfs_mountpoint=${old_rw_overlay_mountpoint}
new_rw_overlay_mountpoint=${old_ro_rootfs_mountpoint}

boot_mountpoint="/boot/firmware"

rootfstype=$(findmnt -n -o fstype /)

if [ "${rootfstype}" = "overlay" ]; then
    # first time through the script when upgrading "normal" overlayfs system


    ## in place major upgrade:
    ## - disable overlay,
    ## - schedule script to run on the underlay using systemd, and
    ## - reboot
    
    mount -o remount,rw ${old_ro_rootfs_mountpoint}

    ###############################
    # Schedule script to run on boot into underlay
    ###############################
    cat <<EOF > ${old_ro_rootfs_mountpoint}/etc/systemd/system/do_major_upgrade.service
[Unit]
Description=Run do-major-upgrade on boot
After=network.target

[Service]
Type=simple
ExecStart=${script_path}
StandardOutput=append:${script_dir}/major_upgrade_final.log
StandardError=append:${script_dir}/major_upgrade_final.log

Restart=on-failure

[Install]
WantedBy=multi-user.target
EOF
    # enable systemd job in read-only rootfs
    overlayroot-chroot systemctl enable do_major_upgrade.service
    # disable jaiabot service in read-only rootfs to reduce startup churn
    overlayroot-chroot systemctl disable jaiabot.service
    # disable JCU so that UI doesn't refresh until upgrade is complete
    overlayroot-chroot systemctl disable jaiabot_goby_liaison_standalone.service
    # Backup/remove rootfs overlayconf file (so system boots into underlay)
    overlayroot-chroot mv /etc/overlayroot.conf /etc/overlayroot.conf.bak
    # reboot into underlay
    systemctl reboot --force
    exit 0
else
    ## second time through when booted on formerly read-only underlay OR
    ## for AWS AMI
    mount $(blkid -L overlay) ${old_rw_overlay_mountpoint}
fi


###############################
# Extract new rootfs
###############################

# delete everything that was on this old overlay
rm -rf ${new_ro_rootfs_mountpoint}/*
sync ${new_ro_rootfs_mountpoint}

# untar the new rootfs
tar --xattrs --xattrs-include="*" -xz -f {{ rootfs_filename }} -C ${new_ro_rootfs_mountpoint} --warning=no-timestamp

if [ "{{ upgrade_type }}" = "embedded" ]; then
    ## in place major upgrade
    
    ###############################
    # Restore reuse files
    ###############################

    rsync -aP {{ upgrade_reuse_dir }}/ ${new_ro_rootfs_mountpoint}/

    ###############################
    # Write new boot dir
    ###############################

    mount -o remount,rw ${boot_mountpoint}
    rsync -a --delete {{ upgrade_boot_dir }}/ ${boot_mountpoint}/ 
fi
    
###############################
# Swap rootfs/overlay labels     
###############################
ROOTFS_DEVICE=$(blkid -L rootfs)
FS_TYPE=$(blkid -o value -s TYPE "$ROOTFS_DEVICE")
echo "Detected filesystem: $FS_TYPE on $ROOTFS_DEVICE"

case "$FS_TYPE" in
    ext4)
        e2label "$ROOTFS_DEVICE" overlay
        ;;
    btrfs)
        ROOTFS_MOUNTPOINT=$(findmnt -n -o TARGET $ROOTFS_DEVICE)
        btrfs filesystem label "$ROOTFS_MOUNTPOINT" overlay
        ;;
    *)
        echo "Error: Unsupported filesystem type '$FS_TYPE'"
        exit 1
        ;;
esac

btrfs filesystem label ${new_ro_rootfs_mountpoint} "rootfs"

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

if [ "{{ upgrade_type }}" = "embedded" ]; then
    rm -rf {{ upgrade_tmp_dir }}
    rm -rf {{ upgrade_reuse_dir }}
    rm -rf {{ upgrade_boot_dir }}
    rm {{ upgrade_dir }}/*.boot.tar.gz
    rm {{ upgrade_dir }}/*.rootfs.tar.gz
    rm {{ upgrade_dir }}/do-major-upgrade.sh
    rm {{ upgrade_dir }}/fleet*.cfg
    mv {{ upgrade_dir_base }}/version.txt {{ upgrade_dir }}
fi

###############################
# Reboot    
###############################

if [ -e /boot/grub/grub.cfg ]; then
    ## AWS AMI or VirtualBox VM
    
    # Pivot to new root
    mount --make-rprivate /
    mkdir -p ${new_ro_rootfs_mountpoint}/oldroot    
    pivot_root ${new_ro_rootfs_mountpoint} ${new_ro_rootfs_mountpoint}/oldroot

    for i in dev proc sys run boot/firmware var/log; do mount --move /oldroot/$i /$i; done

    # move any mounts under /oldmount as well
    mount | awk '{print $3}' | grep "^/oldroot/" | sort -r | while read -r mountpoint; do
        new_mount="${mountpoint//oldroot\//}"
        mount --move "$mountpoint" "$new_mount"
    done

    # install grub
    root_partition=$(findmnt -n -o SOURCE /)
    root_device=$(lsblk -no PKNAME "$root_partition")

    mount -o remount,rw $boot_mountpoint
    mount -o bind $boot_mountpoint /boot/efi

    if ! grub-install "/dev/$root_device" --no-uefi-secure-boot --removable; then
        # older VMs use BIOS
        apt-get -y remove grub-efi-amd64
        apt-get -y install grub-pc
        grub-install "/dev/$root_device" --target=i386-pc
    fi

    # update grub
    if ! which cloud-init || [ "$(cloud-init query cloud_id)" = "nocloud" ] || [ "$(cloud-init query cloud_id)" = "" ]; then
        # VirtualBox cloud init source
        sed -i 's|\(GRUB_CMDLINE_LINUX_DEFAULT=".*\)"|\1 ds=nocloud\\;s=file:///etc/jaiabot/init/ network-config=disabled"|' /etc/default/grub
    elif [ "$(cloud-init query cloud_id)" = "aws" ]; then       
        # Enable serial console for AWS debugging
        sed -i 's/GRUB_CMDLINE_LINUX="\(.*\)"/GRUB_CMDLINE_LINUX="\1 console=ttyS0"/' /etc/default/grub
    fi
       
    update-grub
fi

if [ "{{ upgrade_type }}" = "embedded" ]; then
    systemctl reboot --force
fi

