#!/bin/bash -e

script_dir=$(cd "$(dirname "$0")" && pwd)

if [ -f /etc/default/rdk-miniboot-update ]; then
   . /etc/default/rdk-miniboot-update
fi

LOCAL_MODE=0
if [ -n "$FIRMWARE_ROOT" ]; then
   # Provided by environment
   true
elif [ -d /lib/firmware/rdk/miniboot ]; then
   # Default firmware root exists
   FIRMWARE_ROOT=/lib/firmware/rdk/miniboot
fi

FIRMWARE_RELEASE_STATUS=${FIRMWARE_RELEASE_STATUS:-default}
FIRMWARE_IMAGE_DIR=${FIRMWARE_IMAGE_DIR:-${FIRMWARE_ROOT}/${FIRMWARE_RELEASE_STATUS}}

EXIT_SUCCESS=0
EXIT_FAILED=1

die() {
   echo "$@" >&2
   exit ${EXIT_FAILED}
}

get_uboot_date_iso() {
    local img="$1"
    local miniboot_version

    miniboot_version=$(strings "$img" | grep -oPm1 'U-Boot 2022\.10.*\(\K[^)]+' | sed 's/ - / /')

    if [[ -n "$miniboot_version" ]]; then
        date -d "$miniboot_version" +%s 2>/dev/null || echo "0"
    else
        echo "0"
    fi
}

version_check() {
    local nand_version
    local img_version
    local confirm

    nand_version=$(get_uboot_date_iso /dev/mtd0)
    img_version=$(get_uboot_date_iso "${MINIBOOT_UPDATE_IMAGE}")

    echo "NAND miniboot version timestamp: $nand_version ($(date -d "@$nand_version" '+%Y-%m-%d %H:%M:%S'))"
    echo "IMG miniboot version timestamp: $img_version ($(date -d "@$img_version" '+%Y-%m-%d %H:%M:%S'))"

    if [[ -z "$nand_version" || -z "$img_version" ]]; then
        echo "Failed to parse version info"
        return 1
    fi

    if (( img_version > nand_version )); then
        echo "Image is newer, update needed"
        return 0
    else
        echo "Image is not newer(⚠️  Current version may have hardware issues.)"
        echo "Continue update anyway? [y/N] (5s to respond): "
        if read -t 5 -r confirm; then
            if [[ "$confirm" =~ ^[Yy]$ ]]; then
                echo "Proceeding with update..."
                return 0
            else
                echo "Update aborted by user."
                return 1
            fi
        else
            echo "No input detected. Update aborted."
            return 1
        fi
    fi
}

wrong_version_check() {
    local wrong_nand_version=1747897291
    local nand_version

    nand_version=$(get_uboot_date_iso /dev/mtd0)

    echo "NAND miniboot version timestamp: $nand_version ($(date -d "@$nand_version" '+%Y-%m-%d %H:%M:%S'))"

    if [[ "$nand_version" != "$wrong_nand_version" ]]; then
        echo "nand_version (timestamp): $nand_version is not the known wrong version, skipping auto update"
        return 1
    fi
}

applyUpdate()
{
    [ -n "${MINIBOOT_UPDATE_IMAGE}" ] || die "No Miniboot image specified"
    if [ -n "${MINIBOOT_UPDATE_IMAGE}" ]; then
        [ -f "${MINIBOOT_UPDATE_IMAGE}" ] || die "Miniboot image \"${MINIBOOT_UPDATE_IMAGE}\" not found"
    fi

    version_check

    # Get image dile size
    image_size=`du -b ${MINIBOOT_UPDATE_IMAGE} | awk '{print $1}'`

    # Get partition size
    partition_size=0x`cat /proc/mtd | grep mtd0 | awk '{print $2}'`
    partition_size_1=0x`cat /proc/mtd | grep mtd1 | awk '{print $2}'`

    # The image file size needs to be greater than or equal to the partition size
    _part_size=`printf %d ${partition_size}`
    _part_size_1=`printf %d ${partition_size_1}`

    echo miniboot=${_part_size} ubootenv=${_part_size_1} img=${image_size}

    if [ ${_part_size} -lt ${image_size} ]; then
        split -b ${_part_size} ${MINIBOOT_UPDATE_IMAGE} ${MINIBOOT_UPDATE_IMAGE}
        image_size_ab=`du -b ${MINIBOOT_UPDATE_IMAGE}ab | awk '{print $1}'`
        echo image_size_ab=${image_size_ab}
        if [ ${_part_size_1} -eq ${image_size_ab} ]; then
            echo "Start Update Miniboot("${MINIBOOT_UPDATE_IMAGE}")..."
            echo "Please do not interrupt, restart and shut down during the upgrade process."

            flash_opts=""
            if [ "${SILENT_UPDATE}" = 1 ];then
            flash_opts="-q"
            fi

            flash_erase ${flash_opts} /dev/mtd0 0 0
            if [ $? -ne 0 ]; then
            die "flash_erase /dev/mtd0 failed, dont reboot, please retry."
            fi
            nandwrite ${flash_opts} -p -N /dev/mtd0 ${MINIBOOT_UPDATE_IMAGE}aa
            if [ $? -ne 0 ]; then
            die "nandwrite /dev/mtd0 failed, dont reboot, please retry."
            fi

            flash_erase ${flash_opts} /dev/mtd1 0 0
            if [ $? -ne 0 ]; then
            die "flash_erase /dev/mtd1 failed, dont reboot, please retry."
            fi
            nandwrite ${flash_opts} -p -N /dev/mtd1 ${MINIBOOT_UPDATE_IMAGE}ab
            if [ $? -ne 0 ]; then
            die "nandwrite /dev/mtd1 failed, dont reboot, please retry."
            fi

            sync
            echo "=============================================="
            echo "Update Miniboot Done.Please reboot to apply the update"
            echo "=============================================="
        else
            echo "The format of the miniboot image is incorrect."
            echo "Please use the miniboot image file officially provided."
            exit ${EXIT_FAILED}
        fi
        rm -rf ${MINIBOOT_UPDATE_IMAGE}aa
        rm -rf ${MINIBOOT_UPDATE_IMAGE}ab
    else
        echo "Start Update Miniboot("${MINIBOOT_UPDATE_IMAGE}")..."
        echo "Please do not interrupt, restart and shut down during the upgrade process."

        flash_opts=""
        if [ "${SILENT_UPDATE}" = 1 ];then
            flash_opts="-q"
        fi

        flash_erase ${flash_opts} /dev/mtd0 0 0
        if [ $? -ne 0 ]; then
            die "flash_erase /dev/mtd0 failed, dont reboot, please retry."
        fi
        nandwrite ${flash_opts} -p -N /dev/mtd0 ${MINIBOOT_UPDATE_IMAGE}
        if [ $? -ne 0 ]; then
            die "nandwrite /dev/mtd0 failed, dont reboot, please retry."
        fi

        sync
        echo "=============================================="
        echo "Update Miniboot Done.Please reboot to apply the update"
        echo "=============================================="
    fi

    exit 0
}

# Find latest applicable update version
MINIBOOT_UPDATE_IMAGE=""
MINIBOOT_UPDATE_VERSION=0
getMinibootUpdateVersion()
{
   MINIBOOT_UPDATE_VERSION=0
   match=".*/disk_nand_minimum_boot_[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9].img"
   latest="$(find "${FIRMWARE_IMAGE_DIR}/" -maxdepth 1 -type f -follow -regex "${match}" | sort -r | head -n1)"
   if [ -f "${latest}" ]; then
      MINIBOOT_UPDATE_VERSION=$(basename "${latest}" | awk -F_ '{print $NF}' | sed 's/\.img//')
      MINIBOOT_UPDATE_IMAGE="${latest}"
   fi
}

checkDependencies()
{
    [ "$(id -u)" = "0" ] || die "* Must be run as root - try 'sudo rdk-miniboot-update'"

    if ! command -v flash_erase > /dev/null; then
        die "flash_erase not found. Try re-installing the mtd-utils package"
    fi

    if ! command -v nandwrite > /dev/null; then
        die "nandwrite not found. Try installing the mtd-utils package."
    fi

    if [ ! -f "/proc/mtd" ]; then
        die "NAND flash partition file \"/proc/mtd\" does not exist."
    fi

    if [ ! -c "/dev/mtd0" ]; then
        die "NAND flash partition file \"/dev/mtd0\"does not exist."
    fi
}

usage() {
cat <<EOF
rdk-miniboot-update [options]... [FILE]

Miniboot update tool for the RDK.

Options:
    -f Install the given file instead of the latest applicable update
    -h Display help text and exit
    -l Returns the full path to the latest available miniboot image file according
        to the FIRMWARE_RELEASE_STATUS and FIRMWARE_IMAGE_DIR settings.
    -s Do not display flash progress messages

Environment:
Environment variables should be defined in /etc/default/rdk-miniboot-update

FIRMWARE_RELEASE_STATUS

Specifies the release status of the firmware to apply.

EOF
  exit ${EXIT_SUCCESS}
}

SILENT_UPDATE=0
while getopts ahlf:s option; do
   case "${option}" in
   f) MINIBOOT_UPDATE_IMAGE="${OPTARG}"
      ;;
   l)
      getMinibootUpdateVersion
      echo "${MINIBOOT_UPDATE_IMAGE}"
      exit 0
      ;;
   s) SILENT_UPDATE=1
      ;;
   h) usage
      ;;
   a)
      wrong_version_check
      ;;
   *) echo "Unknown argument \"${option}\""
      usage
      ;;
   esac
done

checkDependencies
if [ -z "${MINIBOOT_UPDATE_IMAGE}" ];then
    getMinibootUpdateVersion
fi
applyUpdate
