#!/bin/bash

set -e

IMAGE="Armbian-unofficial_26.02.0-trunk_Sk-am62-bq03_trixie_vendor_6.12.57.img"
DEVICE="/dev/mmcblk0"
LOG_FILE="/tmp/flash_emmc.log"
SERIAL_DEV="/dev/ttyS2"

exec > >(tee -a ${LOG_FILE} ${SERIAL_DEV}) 2>&1

echo "========== EMMC FLASH START =========="
date

# -----------------------------
# 1. 基本检查
# -----------------------------
echo "[INFO] Checking image file..."

if [ ! -f "${IMAGE}" ]; then
    echo "[ERROR] Image file not found: ${IMAGE}"
    exit 1
fi

echo "[INFO] Checking target device..."

if [ ! -b "${DEVICE}" ]; then
    echo "[ERROR] Target device not found: ${DEVICE}"
    exit 1
fi

# 防止烧错设备（简单保护）
ROOT_DEV=$(mount | grep " / " | awk '{print $1}')
if [[ "${ROOT_DEV}" == *"${DEVICE}"* ]]; then
    echo "[ERROR] ${DEVICE} is current root device! Abort."
    exit 1
fi

# -----------------------------
# 2. 获取镜像大小
# -----------------------------
IMG_SIZE=$(stat -c%s "${IMAGE}")
echo "[INFO] Image size: $((IMG_SIZE / 1024 / 1024)) MB"

# -----------------------------
# 3. 卸载设备分区
# -----------------------------
echo "[INFO] Unmounting partitions..."
umount ${DEVICE}* 2>/dev/null || true

# -----------------------------
# 4. 开始烧录（带进度）
# -----------------------------
echo "[INFO] Flashing image to ${DEVICE}..."

if command -v pv >/dev/null 2>&1; then
    echo "[INFO] Using pv for progress display"
    pv "${IMAGE}" | dd of=${DEVICE} bs=4M conv=fsync status=none
else
    echo "[WARN] pv not found, fallback to dd progress"
    dd if="${IMAGE}" of=${DEVICE} bs=4M conv=fsync status=progress
fi

if [ $? -ne 0 ]; then
    echo "[ERROR] Flash failed!"
    exit 1
fi

# -----------------------------
# 5. 同步数据
# -----------------------------
echo "[INFO] Syncing data..."
sync

# -----------------------------
# 6. 刷新分区表
# -----------------------------
echo "[INFO] Reloading partition table..."
partprobe ${DEVICE} || true
sleep 2

# -----------------------------
# 7. 校验写入（简单校验）
# -----------------------------
echo "[INFO] Verifying write (first 10MB)..."

dd if=${DEVICE} bs=1M count=10 2>/dev/null | md5sum > /tmp/dev_md5
dd if=${IMAGE} bs=1M count=10 2>/dev/null | md5sum > /tmp/img_md5

if diff /tmp/dev_md5 /tmp/img_md5 >/dev/null; then
    echo "[INFO] Verification OK"
else
    echo "[WARN] Verification mismatch (first 10MB)"
fi

# -----------------------------
# 8. 完成提示
# -----------------------------
echo "========== FLASH COMPLETE =========="
date

echo "[SUCCESS] Image flashed to ${DEVICE}"