#!/bin/bash

SOURCE_DIR="$(dirname $(realpath "${BASH_SOURCE[0]}"))"

if [ "$UID" -ne 0 ]; then
    >&2 echo 'Must be run as root'
    exit 1
fi

POSITIONAL_ARGS=()

while [[ $# -gt 0 ]]; do
    case $1 in
      --no-sync)
          NO_SYNC=1
          shift
          ;;
      --create)
          ONLY_CREATE=1
          shift
          ;;
      -*|--*)
          echo "Unknown option $1"
          exit 1
          ;;
      *)
          POSITIONAL_ARGS+=("$1") # save positional arg
          shift # past argument
          ;;
    esac
done

set -- "${POSITIONAL_ARGS[@]}" # restore positional parameters

if [ -z "$NO_SYNC" ]; then
    echo 'Syncing...'
    umount -l /media/startos/next 2> /dev/null
    umount /media/startos/upper 2> /dev/null
    rm -rf /media/startos/upper /media/startos/next
    mkdir /media/startos/upper
    mount -t tmpfs tmpfs /media/startos/upper
    mkdir -p /media/startos/upper/data /media/startos/upper/work /media/startos/next
    mount -t overlay \
		  -olowerdir=/media/startos/current,upperdir=/media/startos/upper/data,workdir=/media/startos/upper/work \
		  overlay /media/startos/next
fi

if [ -n "$ONLY_CREATE" ]; then
    exit 0
fi

mkdir -p /media/startos/next/run
mkdir -p /media/startos/next/dev
mkdir -p /media/startos/next/sys
mkdir -p /media/startos/next/proc
mkdir -p /media/startos/next/boot
mkdir -p /media/startos/next/media/startos/root
mount --bind /run /media/startos/next/run
mount --bind /tmp /media/startos/next/tmp
mount --bind /dev /media/startos/next/dev
mount -t sysfs sysfs /media/startos/next/sys
mount -t proc proc /media/startos/next/proc
mount --bind /boot /media/startos/next/boot
mount --bind /media/startos/root /media/startos/next/media/startos/root

if mountpoint /sys/firmware/efi/efivars 2>&1 > /dev/null; then
  mount -t efivarfs efivarfs /media/startos/next/sys/firmware/efi/efivars
fi

if [ -z "$*" ]; then
    chroot /media/startos/next
    CHROOT_RES=$?
else
    chroot /media/startos/next "$SHELL" -c "$*"
    CHROOT_RES=$?
fi

if mountpoint /media/startos/next/sys/firmware/efi/efivars 2>&1 > /dev/null; then
  umount /media/startos/next/sys/firmware/efi/efivars 
fi

umount /media/startos/next/run
umount /media/startos/next/tmp
umount /media/startos/next/dev
umount /media/startos/next/sys
umount /media/startos/next/proc
umount /media/startos/next/boot
umount /media/startos/next/media/startos/root

if [ "$CHROOT_RES" -eq 0 ]; then

    if [ -h /media/startos/config/current.rootfs ] && [ -e /media/startos/config/current.rootfs ]; then
        ${SOURCE_DIR}/prune-images $(du -s --bytes /media/startos/next | awk '{print $1}')
    fi

    echo 'Upgrading...'

    rm -f /media/startos/images/next.squashfs
    if ! time mksquashfs /media/startos/next /media/startos/images/next.squashfs -b 4096 -comp gzip; then
        umount -l /media/startos/next
        umount -l /media/startos/upper
        rm -rf /media/startos/upper /media/startos/next
        exit 1
    fi
    hash=$(b3sum /media/startos/images/next.squashfs | head -c 32)
    mv /media/startos/images/next.squashfs /media/startos/images/${hash}.rootfs
    ln -rsf /media/startos/images/${hash}.rootfs /media/startos/config/current.rootfs

    sync

    reboot
fi

umount -l /media/startos/next
umount -l /media/startos/upper
rm -rf /media/startos/upper /media/startos/next