#  Copyright (C) Marvell Inc

# over ride default behaviour

echo "Preparing for installation ... "

# global defines
kernel_addr=0x1100000
fdt_addr=0x2800000
fdt_high=0x28fffff
initrd_addr=0x2900000

kernel_fname="/boot/vmlinuz-5.10.0-18-2-armmp"
initrd_fname="/boot/initrd.img-5.10.0-18-2-armmp"
fdt_fname="/boot/armada-385-ET6448M_4G_Nand.dtb"

if [ "$install_env" = "onie" ]; then
    MACH_FILE="/etc/machine.conf"
else
    MACH_FILE="/host/machine.conf"
fi
# armhf-marvell_et6448m_52x-r0  - Platform = Et6448M
# armhf-nokia_ixs7215_52x-r0 - Platform = Nokia IPD6448M
PLATFORM=`sed -n 's/^onie_platform=\(.*\)/\1/p' $MACH_FILE`
echo "Intalling SONiC from $install_env on Platform $PLATFORM"

# global mount defines
if [ "$PLATFORM" = "armhf-marvell_et6448m_52x-r0" ]; then
    demo_dev=ubi0
    mtd_dev=/dev/$(cat /proc/mtd | grep "SONIC" | grep -o "mtd[0-9]")
    mtd_num=$(echo $mtd_dev | grep -o "[0-9]")

    fdt_fname="/boot/armada-385-ET6448M_4G_Nand.dtb"

    BOOTARGS='setenv bootargs root='$demo_dev' rw rootwait ubi.mtd='$mtd_num' rootfstype=ubifs panic=1 console=ttyS0,115200 ${othbootargs} ${mtdparts} ${linuxargs}'
    UBI_LOAD='run ubi_sonic_boot_mount_ubi; ubifsload $kernel_addr $image_name;ubifsload $fdt_addr $fdt_name; ubifsload $initrd_addr $initrd_name'
    BOOTARGS_OLD='setenv bootargs root='$demo_dev' rw rootwait ubi.mtd='$mtd_num' rootfstype=ubifs panic=1 console=ttyS0,115200 ${othbootargs} ${mtdparts} ${linuxargs_old}'
    UBI_LOAD_OLD='run ubi_sonic_boot_mount_ubi; ubifsload $kernel_addr $image_name_old;ubifsload $fdt_addr $fdt_name_old; ubifsload $initrd_addr $initrd_name_old'
    UBIBOOTCMD='run ubi_sonic_boot_bootargs; run ubi_sonic_boot_load; test -n "$boot_once" && setenv boot_once "" && saveenv; bootz $kernel_addr $initrd_addr $fdt_addr'
    UBIBOOTCMD_OLD='run ubi_sonic_boot_bootargs_old; run ubi_sonic_boot_load_old; test -n "$boot_once" && setenv boot_once "" && saveenv; bootz $kernel_addr $initrd_addr $fdt_addr'
    LINUX_MISC_CMD='apparmor=1 security=apparmor usbcore.autosuspend=-1'

elif [ "$PLATFORM" = "armhf-nokia_ixs7215_52x-r0" ]; then
    demo_dev=sda2

    fdt_fname="/boot/armada-385-ipd6448m.dtb"

    BOOTARGS='setenv bootargs root=/dev/'$demo_dev' rw rootwait rootfstype=ext4 panic=1 console=ttyS0,${baudrate} ${othbootargs} ${mtdparts} ${linuxargs}'
    UBI_LOAD='scsi init; ext4load scsi 0:2 $kernel_addr $image_name; ext4load scsi 0:2 $fdt_addr $fdt_name; ext4load scsi 0:2 $initrd_addr $initrd_name'
    BOOTARGS_OLD='setenv bootargs root=/dev/'$demo_dev' rw rootwait rootfstype=ext4 panic=1 console=ttyS0,${baudrate} ${othbootargs} ${mtdparts} ${linuxargs_old}'
    UBI_LOAD_OLD='scsi init; ext4load scsi 0:2 $kernel_addr $image_name_old; ext4load scsi 0:2 $fdt_addr $fdt_name_old; ext4load scsi 0:2 $initrd_addr $initrd_name_old'
    UBIBOOTCMD='run ubi_sonic_boot_bootargs; run ubi_sonic_boot_load; test -n "$boot_once" && setenv boot_once "" && saveenv; bootz $kernel_addr $initrd_addr $fdt_addr'
    UBIBOOTCMD_OLD='run ubi_sonic_boot_bootargs_old; run ubi_sonic_boot_load_old; test -n "$boot_once" && setenv boot_once "" && saveenv; bootz $kernel_addr $initrd_addr $fdt_addr'
    LINUX_MISC_CMD='apparmor=1 security=apparmor usbcore.autosuspend=-1'

else
    echo "Unsupported Platform $PLAFORM"
    exit 2
fi
demo_mnt=/tmp
FW_ENV_DEFAULT='/dev/mtd0 0x00500000 0x80000 0x100000 8'
UBOOT_FW_DEFAULT=1

prepare_boot_menu() {
    echo "Sync up cache ..."
    sync
    echo "Setting up U-Boot environment..."

    DTB_HAS_ENV_BLK=$(grep uboot-env /proc/mtd | sed -e 's/:.*$//')
    if [ -c "/dev/$DTB_HAS_ENV_BLK" ]; then
        PROC_ENV_FILE=$(find /proc/device-tree/ -name env_size)
        if [ -n "$PROC_ENV_FILE" ]
        then
            UBOOT_ENV_SIZ="0x$(hd $PROC_ENV_FILE | awk 'FNR==1 {print $2 $3 $4 $5}')"
            UBOOT_ENV_ERASE_SIZ="0x$(grep uboot-env /proc/mtd | awk '{print $3}')"
            if [[ -n "$UBOOT_ENV_SIZ" && -n "$UBOOT_ENV_ERASE_SIZ" ]]
            then
                # Env info from DTB 
                FW_ENV_DTB="/dev/$DTB_HAS_ENV_BLK 0x00000000 $UBOOT_ENV_SIZ $UBOOT_ENV_ERASE_SIZ"
            fi
        fi
    fi
    if [ -n "$FW_ENV_DTB" ]
    then
                echo $FW_ENV_DTB > /etc/fw_env.config
                echo "Found uboot env offset in device tree"
                UBOOT_ENV_CRC=$(fw_printenv 2>&1 | grep -c 'Warning: Bad CRC') || true
                if [ "$UBOOT_ENV_CRC" -eq 0 ]
                then
                    echo "Uboot env offset in device tree is valid"
                    UBOOT_FW_DEFAULT=0
                    fdt_fname="/boot/armada-385-ET6448M_4G_Nand_r2019.dtb"
                else
                    echo "Uboot env offset in device tree is NOT valid"
                    UBOOT_FW_DEFAULT=1
                fi
    fi
    if [ $UBOOT_FW_DEFAULT -eq 1 ]
    then
        echo $FW_ENV_DEFAULT > /etc/fw_env.config
        echo "Using pre-configured uboot env"
    fi
    if [ "$PLATFORM" = "armhf-nokia_ixs7215_52x-r0" ]; then
        FW_ENV_DEFAULT='/dev/mtd0 0x00100000 0x10000 0x10000'
        echo $FW_ENV_DEFAULT > /etc/fw_env.config
        echo "Using pre-configured uboot env for armhf-nokia_ixs7215_52x-r0"
    fi

    image_name=${image_dir}${kernel_fname}
    initrd_name=${image_dir}${initrd_fname}
    fdt_name=${image_dir}${fdt_fname}

    if [ "$install_env" = "onie" ]; then
        FW_ARG="-f"
        image_dir_old=""
        image_name_old=""
        initrd_name_old=""
        fdt_name_old=""
        sonic_version_2="None"
    else
        eval running_version="$(cat /proc/cmdline | sed -n 's/^.*loop=\/*image-\(\S\+\)\/.*$/\1/p')"
        current_boot_next=$(fw_printenv -n boot_next || true)
        running_sonic_image=1
        if [ "$current_boot_next" = "run sonic_image_1" ]; then
            # boot_next=run sonic_image_1
            sonic_version_2=$(fw_printenv -n sonic_version_2 || true)
            if test "${sonic_version_2#*$running_version}" != "$sonic_version_2"
            then
                running_sonic_image=2
            else
                running_sonic_image=1
            fi
        else
            # boot_next=run sonic_image_2  
            sonic_version_1=$(fw_printenv -n sonic_version_1 || true)
            if test "${sonic_version_1#*$running_version}" != "$sonic_version_1"
            then
                running_sonic_image=1
            else
                running_sonic_image=2
            fi
        fi   
        
        if [ $running_sonic_image -eq 1 ]; then
            image_dir_old=$(fw_printenv -n image_dir || true)
            image_name_old=$(fw_printenv -n image_name || true)
            initrd_name_old=$(fw_printenv -n initrd_name || true)
            fdt_name_old=$(fw_printenv -n fdt_name || true)
            sonic_version_2=$(fw_printenv -n sonic_version_1 || true)
        else
            image_dir_old=$(fw_printenv -n image_dir_old || true)
            image_name_old=$(fw_printenv -n image_name_old || true)
            initrd_name_old=$(fw_printenv -n initrd_name_old || true)
            fdt_name_old=$(fw_printenv -n fdt_name_old || true)
            sonic_version_2=$(fw_printenv -n sonic_version_2 || true)
        fi
    fi

    # Set boot variables
    fw_setenv ${FW_ARG} image_dir $image_dir > /dev/null
    fw_setenv ${FW_ARG} image_name $image_name > /dev/null
    fw_setenv ${FW_ARG} initrd_name $initrd_name > /dev/null
    fw_setenv ${FW_ARG} fdt_name $fdt_name > /dev/null
    fw_setenv ${FW_ARG} sonic_version_1 $demo_volume_revision_label > /dev/null
    fw_setenv ${FW_ARG} image_dir_old $image_dir_old > /dev/null
    fw_setenv ${FW_ARG} image_name_old $image_name_old > /dev/null
    fw_setenv ${FW_ARG} initrd_name_old $initrd_name_old > /dev/null
    fw_setenv ${FW_ARG} fdt_name_old $fdt_name_old > /dev/null
    fw_setenv ${FW_ARG} sonic_version_2 $sonic_version_2 > /dev/null
    BOOT1='echo " > Boot1: $sonic_version_1 - run sonic_image_1";echo;'
    BOOT2='echo " > Boot2: $sonic_version_2 - run sonic_image_2";echo;'
    if [ "$PLATFORM" = "armhf-nokia_ixs7215_52x-r0" ]; then
        BOOT3='echo " > Boot3: ONIE - run onie_bootcmd";echo;'
    else
        BOOT3='echo " > Boot3: ONIE - run onie_nand_boot";echo;'
    fi
    BORDER='echo "---------------------------------------------------";echo;'
    fw_setenv ${FW_ARG} print_menu "$BORDER $BOOT1 $BOOT2 $BOOT3 $BORDER" > /dev/null
    fw_setenv ${FW_ARG} linuxargs "net.ifnames=0 loopfstype=squashfs loop=$image_dir/$FILESYSTEM_SQUASHFS systemd.unified_cgroup_hierarchy=0 varlog_size=$VAR_LOG_SIZE loglevel=4 ${extra_cmdline_linux} ${ONIE_PLATFORM_EXTRA_CMDLINE_LINUX}" > /dev/null
    fw_setenv ${FW_ARG} linuxargs_old "net.ifnames=0 loopfstype=squashfs loop=$image_dir_old/$FILESYSTEM_SQUASHFS systemd.unified_cgroup_hierarchy=0 varlog_size=$VAR_LOG_SIZE loglevel=4" > /dev/null

    # Set boot configs
    fw_setenv ${FW_ARG} kernel_addr $kernel_addr > /dev/null
    fw_setenv ${FW_ARG} fdt_addr $fdt_addr > /dev/null
    fw_setenv ${FW_ARG} fdt_high $fdt_high > /dev/null
    fw_setenv ${FW_ARG} initrd_addr $initrd_addr > /dev/null
    fw_setenv ${FW_ARG} mtdids 'nand0=armada-nand' > /dev/null
    if [ $UBOOT_FW_DEFAULT -eq 1 ]
    then
        fw_setenv ${FW_ARG} mtdparts 'mtdparts=armada-nand:10m(U-Boot)ro,20m@10m(ONIE),-(SONIC)' > /dev/null
    else
        fw_setenv ${FW_ARG} mtdparts 'mtdparts=armada-nand:4m(uboot)ro,5m@5m(uboot-env),20m@10m(ONIE),-(SONIC)' > /dev/null
    fi
    if [ "$PLATFORM" = "armhf-nokia_ixs7215_52x-r0" ]; then
        fw_setenv ${FW_ARG} mtdparts 'mtdparts=spi0.0:4m(boot),-(spi-rootfs)' > /dev/null
    fi
    fw_setenv ${FW_ARG} ubi_sonic_boot_mount_ubi 'ubi part SONIC; ubifsmount ubi0' > /dev/null
    fw_setenv ${FW_ARG} ubi_sonic_boot_bootargs "$BOOTARGS" > /dev/null
    fw_setenv ${FW_ARG} ubi_sonic_boot_load "$UBI_LOAD" > /dev/null
    fw_setenv ${FW_ARG} sonic_image_1 "$UBIBOOTCMD" > /dev/null
    fw_setenv ${FW_ARG} ubi_sonic_boot_bootargs_old "$BOOTARGS_OLD" > /dev/null
    fw_setenv ${FW_ARG} ubi_sonic_boot_load_old "$UBI_LOAD_OLD" > /dev/null
    fw_setenv ${FW_ARG} sonic_image_2 "$UBIBOOTCMD_OLD" > /dev/null
    fw_setenv ${FW_ARG} boot_next  'run sonic_image_1'> /dev/null
    fw_setenv ${FW_ARG} bootcmd 'run print_menu; usb start; test -n "$boot_once" && run boot_once; run boot_next' > /dev/null

}

create_ubi_partition() {
    ubidetach /dev/ubi_ctrl -d 0 2>/dev/null || true 

    trap_push "${onie_bin} ubidetach -d 0 /dev/ubi_ctrl || true"
    echo -en "Format mtd partition '$mtd_dev' "
    ubiformat $mtd_dev -y -q || {
        echo "Failed"
    }

    #attaches MTD devices (which describe raw flash) to UBI and creates corresponding UBI devices; ('-m 2' --> mtd2)
    echo -en "ubiattach mtd '$mtd_num' "
    ubiattach /dev/ubi_ctrl -m $mtd_num || {
        echo "Failed"
    }

    #creates UBI volumes on UBI devices
    ubimkvol  /dev/$demo_dev -N $demo_dev -s 3900MiB
}

create_gpt_partition() {
    blk_dev="/dev/sda"
    demo_part=$(sgdisk -p $blk_dev | grep -e "$demo_volume_label" | awk '{print $1}')
    # ONIE partition size 168MB
    onie_part_size=168

    if [ -z "$demo_part" ] ; then
        # Partition Does NOT Exists
        echo "SONIC label [$demo_volume_label] is NOT found in Partition"
        echo "Proceeding to create partition"

        attr_bitmask="0x0"

        sgdisk --new=${demo_part}::+${onie_part_size}MB \
            --attributes=${demo_part}:=:$attr_bitmask \
            --change-name=${demo_part}:$demo_volume_label $blk_dev \
        || {
            echo "Warning: The first trial of creating partition failed, trying the largest aligned available block of sectors on the disk"
            begin=$(sgdisk -F $blk_dev)
            end=$(sgdisk -E $blk_dev)
            sgdisk --new=${demo_part}:$begin:$end \
                --attributes=${demo_part}:=:$attr_bitmask \
                --change-name=${demo_part}:$demo_volume_label $blk_dev
        } || {
            echo "Error: Unable to create partition $demo_part on $blk_dev"
            exit 1
        }
        partprobe || true
    else
        # Partition Exists
        echo "Partition exists $demo_part on $blk_dev"
    fi

    # Make filesystem
    mkfs.ext4 -L $demo_volume_label /dev/$demo_dev
}

create_partition() {
    # Platform speicific partition 
    if [ "$PLATFORM" = "armhf-marvell_et6448m_52x-r0" ]; then
        echo "Doing UBI partition"
        create_ubi_partition
    else
        echo "Doing GPT partition"
        create_gpt_partition
    fi
}

mount_partition() {

    demo_mnt=$(mktemp -d) || {
        echo "Error: Unable to create file sstem mount point"
        exit 1
    }

    echo "Mounting $demo_dev on $demo_mnt "
    if [ "$PLATFORM" = "armhf-marvell_et6448m_52x-r0" ]; then
        trap_push "${onie_bin} umount /dev/ubi0_0|| true"
        mount -t ubifs /dev/ubi0_0 $demo_mnt || {
            echo "mount -t ubifs /dev/ubi0_0 $demo_mnt Failed"
            exit 1
        }
    else
        trap_push "${onie_bin} umount /dev/sda2|| true"
        mount -t ext4 /dev/sda2 $demo_mnt || {
            echo "mount -t ext4 /dev/sda2 $demo_mnt Failed"
            exit 1
        }
    fi
}

bootloader_menu_config() {
    # Update uboot Environment
    prepare_boot_menu
}