#  Copyright (C) Marvell Inc

# over ride default behaviour

echo "Preparing for installation ... "

# global defines
kernel_addr=0x1100000
fdt_addr=0x1000000
initrd_addr=0x2000000
VAR_LOG=512

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

# global mount defines
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]")
demo_mnt=/tmp
FW_ENV_DEFAULT='/dev/mtd0 0x00500000 0x80000 0x100000 8'
UBOOT_FW_DEFAULT=1

# Skip VID Header in UBIFS
BOOTARGS='setenv bootargs root='$demo_dev' rw rootwait ubi.mtd='$mtd_num',8192 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',8192 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'

prepare_boot_menu() {
    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
    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
        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)
    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;'
    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 varlog_size=$VAR_LOG" > /dev/null
    fw_setenv ${FW_ARG} linuxargs_old "net.ifnames=0 loopfstype=squashfs loop=$image_dir_old/$FILESYSTEM_SQUASHFS varlog_size=$VAR_LOG" > /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} 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
    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_partition() {
    # Platform speicific partition 
    create_ubi_partition
}

mount_partition() {

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

    echo "Mounting $demo_dev on $demo_mnt "
    trap_push "${onie_bin} umount /dev/ubi0_0|| true"
    mount -t ubifs /dev/ubi0_0 $demo_mnt || {
        echo "Failed"
    }
}

bootloader_menu_config() {
    # Update uboot Environment
    prepare_boot_menu
}