#!/bin/bash
set -euo pipefail

export PATH="/usr/hobot/bin:${PATH}"

# Base directory of miniboot OTA packages
readonly OTA_MINIBOOT_BASE_DIR="/usr/lib/firmware/rdk/miniboot/stable"

# Build type: release | debug (default: release)
BUILD_TYPE="release"

REBOOT_CHOICE=""
CONFIRM_CHOICE=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --reboot)
            if [[ $# -lt 2 ]]; then
                echo "Error: --reboot requires an argument: y|n"
                exit 1
            fi
            REBOOT_CHOICE="$2"
            shift 2
            ;;
        --confirm)
            if [[ $# -lt 2 ]]; then
                echo "Error: --confirm requires an argument: y|n"
                exit 1
            fi
            CONFIRM_CHOICE="$2"
            shift 2
            ;;
        --build|--type)
            if [[ $# -lt 2 ]]; then
                echo "Error: $1 requires an argument: release|debug"
                exit 1
            fi
            BUILD_TYPE="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            echo "Usage: $0 [--build release|debug] [--reboot y|n] [--confirm y|n]"
            exit 1
            ;;
    esac
done

# Validate build type (omit => default release)
case "${BUILD_TYPE}" in
    release|debug)
        ;;
    *)
        echo "Invalid build type: ${BUILD_TYPE}"
        echo "Use: --build release|debug (default: release)"
        exit 1
        ;;
esac

# Select package directory according to build type
readonly OTA_MINIBOOT_PACKAGE_DIR="${OTA_MINIBOOT_BASE_DIR}/${BUILD_TYPE}/img_packages"

REBOOT_FLAG=""

if [[ -z "${REBOOT_CHOICE}" ]]; then
    # --reboot was not specified, prompting the user.
    read -r -p "Do you want to reboot now? (y/n): " USER_INPUT
    case "$USER_INPUT" in
        y|Y)
            REBOOT_FLAG=""   # Restart immediately: without -n
            ;;
        n|N)
            REBOOT_FLAG="-n" # To avoid restarting: add -n
            ;;
        *)
            echo "Invalid input. Please run the script again with valid choice."
            exit 1
            ;;
    esac
else
    # The user specified --reboot
    case "$REBOOT_CHOICE" in
        y|Y)
            REBOOT_FLAG=""   # Restart immediately
            ;;
        n|N)
            REBOOT_FLAG="-n" # No restart
            ;;
        *)
            echo "Invalid --reboot value: ${REBOOT_CHOICE}"
            echo "Use: --reboot y|n"
            exit 1
            ;;
    esac
fi

# ---------------------------------------------------------------------------
# Upgrade strategy:
#   NOR  : whole-flash burn of disk/miniboot_flash[_nose].img via dd to
#          /dev/block/platform/by-name/hb_vspiflash. miniboot_flash.img is
#          already laid out by FPT order with primary+bak and both AB slots
#          populated, so one dd covers everything.
#          hobot_vspi (secure) rejects raw chunked writes (HSM auth fails),
#          so we temporarily swap to hobot_vspi_debug, burn, swap back.
#   eMMC : per-partition dd (block device, GPT layout stable, current slot).
#   Image selection by board cmdline:
#     hobotboot.secureboot=1 -> miniboot_flash.img
#     otherwise              -> miniboot_flash_nose.img
# ---------------------------------------------------------------------------

# Refresh /dev/block/platform/by-name MTD symlinks after a driver swap.
# /etc/init.d/by-name.sh uses 'ln -s' (no -f), so stale links from the
# previous module instance are removed first.
_refresh_mtd_bynames() {
    local d="/dev/block/platform/by-name"
    local link
    for link in "${d}"/*; do
        [[ -L "${link}" ]] || continue
        local tgt; tgt=$(readlink "${link}")
        [[ "${tgt}" == /dev/mtd* ]] && rm -f "${link}"
    done
    /etc/init.d/by-name.sh mtd
}

# Whole-NOR burn. Returns 0 on success.
# Idempotent driver swap: paired deb preinst normally swaps hobot_vspi ->
# hobot_vspi_debug before us, but this script is also usable standalone, so
# we still handle the swap here. If we're already on debug, it's a no-op.
# Leave the debug driver loaded after dd — switching back can fail if a
# paired linux-image upgrade has already removed the old kernel's modules.
_nor_whole_burn() {
    local img="$1"
    local dev="/dev/block/platform/by-name/hb_vspiflash"

    if [[ ! -f "${img}" ]]; then
        echo "  FAIL: image not found: ${img}" >&2
        return 1
    fi

    if lsmod | awk '{print $1}' | grep -qx "hobot_vspi"; then
        echo "  Switching driver: hobot_vspi -> hobot_vspi_debug"
        if ! rmmod hobot_vspi 2>&1; then
            echo "  FAIL: rmmod hobot_vspi" >&2
            return 1
        fi
        if ! modprobe hobot_vspi_debug 2>&1; then
            echo "  FAIL: modprobe hobot_vspi_debug" >&2
            modprobe hobot_vspi 2>/dev/null
            return 1
        fi
        _refresh_mtd_bynames
    elif ! lsmod | awk '{print $1}' | grep -qx "hobot_vspi_debug"; then
        echo "  FAIL: neither hobot_vspi nor hobot_vspi_debug is loaded" >&2
        return 1
    fi

    # Whole-flash burn: bs=2M is required (driver IPC limit), 64KB-aligned.
    echo "  dd ${img##*/} -> hb_vspiflash (bs=2M, ~3.5min for 22MB) ..."
    if ! dd if="${img}" of="${dev}" bs=2M status=progress 2>&1; then
        sync
        echo "  FAIL: dd error" >&2
        return 1
    fi
    sync
    echo "  OK"
    return 0
}

direct_flash_upgrade() {
    local img_dir="$1"
    local part_by_name="/dev/block/platform/by-name"
    local emmc_ab_parts="acore_cfg bl31 optee uboot"

    echo ""
    echo "========================================================"
    echo "  !!!  WARNING / 警告  !!!"
    echo "========================================================"
    echo ""
    echo "  [EN] Upgrading via direct flashing."
    echo "       This is a LOW-LEVEL operation."
    echo ""
    echo "       - Back up all important data BEFORE proceeding."
    echo "       - Do NOT power off, reboot, or perform ANY"
    echo "         operation on the device during the upgrade."
    echo "       - Interrupting the process MAY BRICK the device."
    echo ""
    echo "  [CN] 直接烧写模式升级。"
    echo "       这是底层操作，请注意："
    echo ""
    echo "       - 升级前请做好重要数据的备份。"
    echo "       - 升级过程中切勿断电、重启或对设备进行任何操作。"
    echo "       - 升级中断可能导致设备变砖，无法恢复。"
    echo ""
    echo "========================================================"
    echo ""
    local confirm
    if [[ -z "${CONFIRM_CHOICE}" ]]; then
        read -r -p "Confirm to proceed? / 确认继续? (y/N): " confirm
    else
        confirm="${CONFIRM_CHOICE}"
    fi
    case "${confirm}" in
        y|Y) ;;
        *) echo "Aborted. / 已取消。"; exit 1 ;;
    esac
    echo ""

    # Board secure mode (cmdline is authoritative; driver may be swapped at runtime)
    local board_secure=0
    grep -q "hobotboot.secureboot=1" /proc/cmdline 2>/dev/null && board_secure=1

    # OHP lifecycle blocks on-device NOR rewrites even via the debug driver.
    if /usr/hobot/bin/provision_tool --get-lifecycle 2>/dev/null | grep -q OHP; then
        echo "" >&2
        echo "Error: device is in OHP lifecycle." >&2
        echo "       Online (on-device) miniboot upgrade is NOT supported." >&2
        echo "       Please reflash the whole image via the factory tool." >&2
        echo "" >&2
        echo "错误：设备已进入 OHP 阶段，禁止在线升级 miniboot。" >&2
        echo "      请使用 factory 工具整片烧录镜像。" >&2
        exit 1
    fi

    # Boot slot for eMMC AB partitions
    local slot_output slot
    slot_output=$(ota_tool -g 2>&1)
    if echo "${slot_output}" | grep -qi "current slot is:A"; then
        slot="a"
    elif echo "${slot_output}" | grep -qi "current slot is:B"; then
        slot="b"
    else
        echo "Error: cannot determine current slot." >&2
        exit 1
    fi

    # NOR whole-flash image by board secure mode
    local nor_img_name
    if [[ "${board_secure}" == "1" ]]; then
        nor_img_name="miniboot_flash.img"
    else
        nor_img_name="miniboot_flash_nose.img"
    fi
    local nor_img="${img_dir}/disk/${nor_img_name}"

    echo "  Slot: ${slot^^}  |  BoardSecure: ${board_secure}  |  NOR image: ${nor_img_name}"
    echo ""

    local failed=0

    # ---- NOR whole-disk burn ----
    echo "  --- NOR whole-disk burn ---"
    if ! _nor_whole_burn "${nor_img}"; then
        failed=$((failed+1))
    fi

    # ---- eMMC/UFS per-partition burn ----
    echo ""
    echo "  --- eMMC/UFS per-partition burn ---"
    _emmc_flash() {
        local part="$1" img="$2" dev="$3"
        printf "  %-35s -> %-25s ... " "${img##*/}" "${dev##*/}"
        if [[ ! -e "${dev}" ]]; then echo "FAIL (device not found)"; failed=$((failed+1)); return; fi
        if dd if="${img}" of="${dev}" bs=4M 2>/dev/null && sync; then echo "OK"
        else echo "FAIL"; failed=$((failed+1)); fi
    }
    for part in ${emmc_ab_parts}; do
        local img="${img_dir}/${part}.img"
        if [[ ! -f "${img}" ]]; then echo "  SKIP (no image): ${part}"; continue; fi
        _emmc_flash "${part}" "${img}" "${part_by_name}/${part}_${slot}"
    done

    if [[ "${failed}" -gt 0 ]]; then
        echo ""
        echo "Error: ${failed} step(s) failed to flash." >&2
        echo "错误：${failed} 个步骤烧写失败。" >&2
        exit 1
    fi
    echo ""
    echo "Direct flash upgrade completed."
}
# ---------------------------------------------------------------------------

direct_flash_upgrade "${OTA_MINIBOOT_PACKAGE_DIR}"

echo "Rdk miniboot update finished successfully. (build=${BUILD_TYPE})"

if [[ -z "${REBOOT_FLAG}" ]]; then
    reboot
fi
