#!/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: direct dd-based partition flash
# ---------------------------------------------------------------------------
direct_flash_upgrade() {
    local img_dir="$1"
    local part_by_name="/dev/block/platform/by-name"
    local ab_parts="spl MCU acore_cfg bl31 optee uboot"
    local fixed_parts="HSM_FW HSM_RCA keyimage SBL"

    echo ""
    echo "========================================================"
    echo "  !!!  WARNING / 警告  !!!"
    echo "========================================================"
    echo ""
    echo "  [EN] Upgrading via direct partition 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 ""

    # Detect secure board
    local secure=0
    [[ -d "${img_dir}/ota_signed_img" ]] && secure=1

    # Detect OHP (same logic as ota_process/ota_update.c)
    local is_ohp=0
    /usr/hobot/bin/provision_tool --get-lifecycle 2>/dev/null | grep -q OHP && is_ohp=1

    # Get current boot slot
    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
    echo "  Slot: ${slot^^}  |  Secure: ${secure}  |  OHP: ${is_ohp}"
    echo ""

    # Resolve image path for a given partition name
    _find_img() {
        local part="$1"
        [[ "${part}" == "keyimage" && "${is_ohp}" == "1" ]] && part="keyimage_ohp"
        local f=""
        if [[ "${secure}" == "1" ]]; then
            f=$(ls "${img_dir}/ota_signed_img/${part}_signed.img" 2>/dev/null || \
                ls "${img_dir}/ota_signed_img/${part}"*"_signed.img" 2>/dev/null | head -1 || true)
            [[ -n "${f}" ]] && echo "${f}" && return
        fi
        f=$(ls "${img_dir}/${part}.img" 2>/dev/null || \
            ls "${img_dir}/${part}_"*.img 2>/dev/null | grep -v signed | head -1 || true)
        echo "${f}"
    }

    local failed=0
    _flash() {
        local part="$1" img="$2" dev="$3"
        printf "  %-35s -> %-20s ... " "${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 ${fixed_parts}; do
        local img; img=$(_find_img "${part}")
        [[ -z "${img}" ]] && { echo "  SKIP (no image): ${part}"; continue; }
        _flash "${part}" "${img}" "${part_by_name}/${part}"
    done
    for part in ${ab_parts}; do
        local img; img=$(_find_img "${part}")
        [[ -z "${img}" ]] && { echo "  SKIP (no image): ${part}"; continue; }
        _flash "${part}" "${img}" "${part_by_name}/${part}_${slot}"
    done

    if [[ "${failed}" -gt 0 ]]; then
        echo ""
        echo "Error: ${failed} partition(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
