#  Copyright (C) Marvell Inc

# over ride default behaviour

echo "Preparing for installation ... "

IS_MASS=0
# global defines
kernel_addr=0x1100000
fdt_addr=0x1000000
fit_addr=0x1000000
initrd_addr=0x2000000
VAR_LOG=512
DISK_SIZE=9

kernel_fname="/boot/vmlinuz-4.9.0-9-2-arm64"
initrd_fname="/boot/initrd.img-4.9.0-9-2-arm64"
fdt_fname="/boot/armada-7020-comexpress.dtb"
fit_fname="/boot/sonic_arm64.fit"

# global mount defines
#demo_dev=/dev/sda1
demo_mnt=/tmp
#mtd_dev=/dev/$(cat /proc/mtd | grep "ENV" | grep -o "mtd[0-9]")
FW_ENV_DEFAULT='/dev/mtd1 0x0 0x10000 0x100000'
UBOOT_FW_DEFAULT=1

# Skip VID Header in UBIFS
LINUX_MISC_CMD='apparmor=1 security=apparmor usbcore.autosuspend=-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 [ $UBOOT_FW_DEFAULT -eq 1 ]
    then
        echo $FW_ENV_DEFAULT > /etc/fw_env.config
        echo "Using pre-configured uboot env"
    fi
    image_name=${image_dir}${kernel_fname}
    initrd_name=${image_dir}${initrd_fname}
    fdt_name=${image_dir}${fdt_fname}
    fit_name=${image_dir}${fit_fname}

    if [ "$install_env" = "onie" ]; then
        FW_ARG="-f"
        image_dir_old=""
        image_name_old=""
        initrd_name_old=""
        fdt_name_old=""
        fit_name_old=""
        sonic_version_2="None"
    else
        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)
        fit_name_old=$(fw_printenv -n fit_name || true)
        sonic_version_2=$(fw_printenv -n sonic_version_1 || true)
        if [ -z "$demo_dev" ] 
        then 
            get_install_device
            if [ $? -ne 0 ]; then
                echo "Error: Unable to detect $blk_dev $demo_dev"
                exit 1
            fi
            if [ ${IS_MASS} -eq 1 ]
            then
                demo_dev=${blk_dev}1
            else
                #demo_dev=$(echo $blk_dev | sed -e 's/\(mmcblk[0-9]\)/\1p/')$demo_part
                demo_dev=/dev/mmcblk0p1
            fi
        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} fit_name $fit_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} fit_name_old $fit_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;'
    BOOT3='echo " > Boot3: ONIE - run onie-nand-boot";echo;'
    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" > /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" > /dev/null
    sonic_bootargs_old='setenv bootargs root='$demo_dev' rw rootwait rootfstype=ext4 panic=1 console=ttyS0,115200 ${othbootargs} ${mtdparts} ${linuxargs_old}'
    fw_setenv ${FW_ARG} sonic_bootargs_old $sonic_bootargs_old > /dev/null || true
    sonic_boot_load_old=$(fw_printenv -n sonic_boot_load || true)
    old_str="_old"
    fw_setenv ${FW_ARG} sonic_boot_load_old $sonic_boot_load_old$old_str > /dev/null || true

    fw_setenv ${FW_ARG} kernel_addr $kernel_addr > /dev/null
    fw_setenv ${FW_ARG} fdt_addr $fdt_addr > /dev/null
    fw_setenv ${FW_ARG} fit_addr $fit_addr > /dev/null
    fw_setenv ${FW_ARG} initrd_addr $initrd_addr > /dev/null
    # Set boot configs
    if [ ${IS_MASS} -eq 1 ]
    then
        #USB_LOAD='ext4load usb 0 $kernel_addr $image_name; ext4load usb 0 $fdt_addr $fdt_name; ext4load usb 0 $initrd_addr $initrd_name'
        USB_LOAD='ext4load usb 0 $fit_addr $fit_name'
        fw_setenv ${FW_ARG} sonic_boot_load $USB_LOAD > /dev/null
    else
        #MMC_LOAD='ext4load mmc 0:'$demo_part' $kernel_addr $image_name; ext4load mmc 0:'$demo_part' $fdt_addr $fdt_name; ext4load mmc 0:'$demo_part' $initrd_addr $initrd_name'
        MMC_LOAD='ext4load mmc 0:'$demo_part' $fit_addr $fit_name'
        fw_setenv ${FW_ARG} sonic_boot_load $MMC_LOAD > /dev/null
    fi
    #SONIC_BOOT_CMD='run sonic_bootargs; run sonic_boot_load; booti $kernel_addr $initrd_addr $fdt_addr'
    SONIC_BOOT_CMD='run sonic_bootargs; run sonic_boot_load; bootm $fit_addr'
    SONIC_BOOT_CMD_OLD='run sonic_bootargs_old; run sonic_boot_load_old; bootm $fit_addr'
    BOOTARGS='setenv bootargs root='$demo_dev' rw rootwait rootfstype=ext4 panic=1 console=ttyS0,115200 ${othbootargs} ${mtdparts} ${linuxargs}'
    fw_setenv ${FW_ARG} sonic_bootargs $BOOTARGS > /dev/null
    fw_setenv ${FW_ARG} sonic_bootcmd $SONIC_BOOT_CMD > /dev/null
    fw_setenv ${FW_ARG} sonic_image_2 $SONIC_BOOT_CMD_OLD > /dev/null
    fw_setenv ${FW_ARG} sonic_image_1 "$SONIC_BOOT_CMD" > /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

}

#Get block device
#Default block device is eMMC, if not look for usb storage
get_install_device()
{
    mass_bus="target0:0:0"
    for i in a b c ; do
        if $(ls -l /sys/block/sd$i/device 2>/dev/null | grep -q "$mass_bus") ; then
            echo "/dev/sd$i"
            blk_dev=/dev/sd$i
            IS_MASS=1
            echo "Selected Mass storage $blk_dev"
            return 0
        fi
    done
    mmc_bus="mmc0:aaaa"
    for i in 0 1 2 ; do
        if $(ls -l /sys/block/mmcblk$i/device 2>/dev/null | grep -q "$mmc_bus") ; then
            echo "/dev/mmcblk$i"
            blk_dev=/dev/mmcblk$i
            echo "Selected mmc $blk_dev"
            return 0
        fi
    done

    echo "ERROR storage not found"
    return 1
}

create_demo_partition() {

    if [ ${IS_MASS} -eq 1 ]
    then
        # USB drive
        parted -s $blk_dev rm 1 || true
        partprobe || true
        parted -s --align optimal $blk_dev unit gb mkpart primary 1 $DISK_SIZE || true
        partprobe || true
    else
        # SD CARD
        parted -s /dev/mmcblk0 rm 1
        partprobe
        parted -s --align optimal /dev/mmcblk0 unit gb mkpart primary 1 $DISK_SIZE
        partprobe
    fi
}

create_partition() {
    get_install_device
    if [ $? -ne 0 ]; then
        echo "Error: Unable to detect $blk_dev $demo_dev"
        exit 1
    fi
    # Platform specific partition 
    create_demo_partition
}

mount_partition() {
    if [ ${IS_MASS} -eq 1 ]
    then
        demo_dev=${blk_dev}1
    else
        #demo_dev=$(echo $blk_dev | sed -e 's/\(mmcblk[0-9]\)/\1p/')$demo_part
        demo_dev=/dev/mmcblk0p1
    fi

    # Make filesystem
    echo "demo label: $demo_volume_label. $demo_dev..."
    mkfs.ext4 -L $demo_volume_label $demo_dev

    mount -t ext4 -o defaults,rw $demo_dev $demo_mnt || {
        echo "Error: Unable to mount $demo_dev on $demo_mnt"
        exit 1
    }
}

bootloader_menu_config() {
    # Update uboot Environment
    prepare_boot_menu
}