#!/bin/bash

# Copyright 2024 NVIDIA Corporation.  All rights reserved.

in_target=/usr/cumulus/bin/in-target

has_two_system_partitions() {
    blkid | grep -q 'LABEL="CL-SYSTEM-2"' && echo "yes"
}

current_boot_partition() {
    local rootfs=$(df / | tail -1 | awk '{ print $1 }')
    blkid $rootfs | sed -e 's/^.*CL-SYSTEM-//' -e 's/[^12].*$//'
}

other_boot_partition() {
    local cbp=$( current_boot_partition )
    if [ $cbp = '1' ]; then
        echo '2'
    else
        echo '1'
    fi
}

storage_device() {
    lsblk -p | grep disk | awk '{ print $1 }'
}

uuid_of_partition() {
    local label="$1"
    blkid | grep "LABEL=\"$label\"" | sed -e 's/^.* UUID="//' -e 's/".*$//'
}

device_of_partition() {
    local label="$1"
    blkid | grep "LABEL=\"$label\"" | sed -e 's/:.*$//'
}

copy_etc() {
    if [ "$copy_etc_arg" = "false" ]; then
        echo "Not copying etc since Copy is false"
        return
    fi
    local mountpoint=$1

    if [ -d "$mountpoint" ]; then
        pushd "$mountpoint"
        mkdir etc.new
        echo "Copying /etc to $mountpoint"
        (cd /etc ; tar cpf - . ) | (cd etc.new ; tar xpf - )
        if [ $? -eq 0 ]; then
            now=$(date +%Y%m%d-%H%M%S)
            # keep the original image,lsb,os release files
            cp -p etc/{image,lsb,os}-release etc.new
            mv etc etc.$now
            mv etc.new etc
        else
            echo "Unable to copy /etc to $mountpoint"
        fi
        popd
    else
        echo "$mountpoint is not a directory"
    fi
}

install_grub_packages() {
    local mountpoint="$1"

    # Ensure that $mountpoint/boot/efi/ exists if /boot/efi/ exists
    if [ -d /boot/efi -a ! -d $mountpoint/boot/efi ]; then
        mkdir -p $mountpoint/boot/efi
    fi

    if [ -d /sys/firmware/efi ]; then
        grub_package=grub-efi-amd64 # ??? arm64
    else
        grub_package=grub-pc
    fi
    echo "Checking for ${grub_package} and installing if necessary"
    $in_target $mountpoint dpkg -l | grep -q "${grub_package} " ||
        $in_target $mountpoint dpkg --install /var/cache/late-debs/${grub_package}*.deb
    # Complete configuration of unpacked but unconfigured packages,
    # such as those that were installed as late-debs.
    # onie-tools is important here, because it produces files in /etc/grub.d/
    $in_target $mountpoint dpkg --configure -a
}

update_grub() {
    local part="$1"
    local grub_installer_defaults=etc/default/grub.d/00-installer-defaults.cfg
    local efi_already_mounted=no

    if [ "$part" != $( current_boot_partition ) ]; then
        local mountpoint="/mnt/CL-SYSTEM-$part"
        mkdir -p $mountpoint
        mount LABEL=CL-SYSTEM-$part $mountpoint
        copy_etc $mountpoint
        mkdir -p $mountpoint/boot/grub
        if [ ! -f $mountpoint/$grub_installer_defaults ]; then
            cp /$grub_installer_defaults $mountpoint/$grub_installer_defaults
        fi
        install_grub_packages $mountpoint
        if [ -d /sys/firmware/efi ]; then
            # workaround for bug #3861745 in generating ONIE entry
            mount LABEL="EFI System" $mountpoint/boot/efi
            if [ $? -eq 32 ]; then
                efi_already_mounted=yes
            fi
        fi
        # Ensure that GRUB_CMDLINE_LINUX in the other partition's
        # /etc/default/grub is non-empty and includes the parameters
        # from /proc/cmdline that are in the current partition's
        # /etc/default/grub .
        if [ ! -f $mountpoint/etc/default/grub ]; then
            cp /etc/default/grub $mountpoint/etc/default/grub
        elif grep -q 'GRUB_CMDLINE_LINUX=""' $mountpoint/etc/default/grub ; then
            grub_cmdline_linux=$( grep GRUB_CMDLINE_LINUX= /etc/default/grub | sed -e 's,",\\",g')
            sed -i -e "s/GRUB_CMDLINE_LINUX=\"\"/$grub_cmdline_linux/" $mountpoint/etc/default/grub
        fi
        grep 'GRUB_CMDLINE_LINUX=""' $mountpoint/etc/default/grub
        $in_target $mountpoint update-grub
        mv /boot/grub/grub.cfg /boot/grub/grub.cfg.old
        cp -p $mountpoint/boot/grub/grub.cfg /boot/grub/grub.cfg
        if [ -d /sys/firmware/efi -a "$efi_already_mounted" != "yes" ]; then
            umount $mountpoint/boot/efi
        fi
        umount $mountpoint
    else
        if [ -d /sys/firmware/efi ]; then
            # workaround for bug #3861745 in generating ONIE entry
            mount LABEL="EFI System" /boot/efi
            if [ $? -eq 32 ]; then
                efi_already_mounted=yes
            fi
        fi
        update-grub
        if [ -d /sys/firmware/efi -a "$efi_already_mounted" != "yes" ]; then
            umount /boot/efi
        fi
    fi
}

if [ "$( has_two_system_partitions )" != "yes" ]; then
    echo "This switch does not have two system partitions."
    exit 1
fi

case "$1" in
    [12])
        partition="$1"
        ;;
    current)
        partition=$(current_boot_partition)
        ;;
    other)
        partition=$(other_boot_partition)
        ;;
    *)
        echo "usage: $0 1|2|current|other [copy-etc]"
        exit 1
        ;;
esac

copy_etc_arg=false
if [ "$2" == "copy-etc" ]; then
    copy_etc_arg=true
fi

echo "Updating GRUB to boot from CL-SYSTEM-$partition"
update_grub $partition

