#!/bin/bash

enter_anti_failure_mode ()
{
    disk="$(grub-probe -t disk /)"
    if [ -e /proc/efi -o -e /sys/firmware/efi ]; then # for UEFI BIOS
        /usr/share/ubuntu/scripts/uefi_helper.sh remove ubuntu
    else # for legacy BIOS
        sfdisk $disk -A 1
    fi
}

leave_anti_failure_mode ()
{
    disk="$(grub-probe -t disk /)"
    if [ -e /proc/efi -o -e /sys/firmware/efi ]; then # for UEFI BIOS
        grub-install
    else # for legacy BIOS
        sfdisk $disk -A 2
    fi
}

early_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        enter_anti_failure_mode
    fi
}

late_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        leave_anti_failure_mode

        device="$(grub-probe -t device /)"
        disk="$(grub-probe -t disk /)"

        # For NVMe or eMMC device, the partition is different, such as 'p1'
        # Get the partition number
        partition_number=${device#$disk}

        # Remove the number in the partition number to get prefix
        partition_prefix=${partition_number//[0-9]/}

        device=${disk}${partition_prefix}

        while ! mount | grep cdrom; do sleep 1; done # Waiting for /cdrom ready

        if [ -e /proc/efi -o -e /sys/firmware/efi ]; then # for UEFI BIOS
            rp=${device}2
        else # for legacy BIOS
            rp=${device}1
        fi
        if mount -rw -o remount ${rp} /cdrom; then
            if [ -f /cdrom/boot/grub/grubenv ]; then
                grub-editenv /cdrom/boot/grub/grubenv set rectype=hdd
                # Unset 'ondiskrec' flag while the whole recovery process is finished. (LP: #1054250)
                if grub-editenv /cdrom/boot/grub/grubenv list | grep ondiskrec; then
                    grub-editenv /cdrom/boot/grub/grubenv unset ondiskrec
                fi
            fi
            mount -r -o remount ${rp} /cdrom
        fi
        rm /usr/share/ubuntu/anti-failure
    fi
}

reboot_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        leave_anti_failure_mode
    fi
    reboot
}

force_reboot_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        leave_anti_failure_mode
    fi
    reboot -f
}

poweroff_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        leave_anti_failure_mode
    fi
    poweroff
}

force_poweroff_task ()
{
    if [ -e /usr/share/ubuntu/anti-failure ]; then
        leave_anti_failure_mode
    fi
    poweroff -f
}

case $1 in
    (early)
        early_task
        ;;
    (late)
        late_task
        ;;
    (reboot)
        reboot_task
        ;;
    (force-reboot)
        force_reboot_task
        ;;
    (poweroff)
        poweroff_task
        ;;
    (force-poweroff)
        force_poweroff_task
        ;;
esac
