#!/bin/bash
#===============================================================================
# 从SD卡操作eMMC备份脚本
#===============================================================================

set -e
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

error_msg() { echo -e "${RED}[ERROR] $1${NC}" >&2; exit 1; }
info_msg() { echo -e "${GREEN}[INFO] $1${NC}"; }
warn_msg() { echo -e "${YELLOW}[WARN] $1${NC}"; }

[[ $EUID -eq 0 ]] || error_msg "需要root权限"

EMMC_DEV="/dev/mmcblk0"
USB_DIR="/mnt/XIN_USB/rock5c"
OUTPUT_FILE="${USB_DIR}/rock5c_emmc_$(date +%Y%m%d_%H%M%S).img"

info_msg "从SD卡备份eMMC (智能压缩版)..."
info_msg "eMMC设备: ${EMMC_DEV}"
info_msg "输出到: ${OUTPUT_FILE}"

# 获取原始eMMC大小
ORIGINAL_SIZE=$(blockdev --getsize64 ${EMMC_DEV})
ORIGINAL_SIZE_GB=$((ORIGINAL_SIZE / 1073741824))
info_msg "原始eMMC大小: ${ORIGINAL_SIZE_GB}GB"

# 检查设备
[[ -b "${EMMC_DEV}" ]] || error_msg "eMMC设备不存在: ${EMMC_DEV}"
[[ -d "${USB_DIR}" ]] || error_msg "USB目录不存在: ${USB_DIR}"

# 检查eMMC没有被挂载
if mount | grep -q mmcblk0; then
    error_msg "eMMC还在使用中！请确保从SD卡启动"
fi

info_msg "✅ eMMC设备可安全操作"

# 步骤1: 检查eMMC分区信息
info_msg "步骤1: 分析eMMC分区..."
parted -s ${EMMC_DEV} unit MB print

# 检查eMMC上的数据大小
info_msg "检查eMMC数据使用量..."
TEMP_MOUNT="/tmp/check_emmc"
mkdir -p "${TEMP_MOUNT}"

# 检查分区是否存在
[[ -b "${EMMC_DEV}p1" ]] || error_msg "未找到boot分区 ${EMMC_DEV}p1"
[[ -b "${EMMC_DEV}p2" ]] || error_msg "未找到root分区 ${EMMC_DEV}p2"

# 临时挂载检查使用量
mount ${EMMC_DEV}p1 "${TEMP_MOUNT}" -o ro
EMMC_BOOT_USED=$(df --block-size=1M "${TEMP_MOUNT}" | tail -1 | awk '{print $3}')
umount "${TEMP_MOUNT}"

mount ${EMMC_DEV}p2 "${TEMP_MOUNT}" -o ro
EMMC_ROOT_USED=$(df --block-size=1M "${TEMP_MOUNT}" | tail -1 | awk '{print $3}')
umount "${TEMP_MOUNT}"

rm -rf "${TEMP_MOUNT}"

TOTAL_USED=$((EMMC_BOOT_USED + EMMC_ROOT_USED))
info_msg "eMMC使用情况: Boot ${EMMC_BOOT_USED}MB + Root ${EMMC_ROOT_USED}MB = 总计 ${TOTAL_USED}MB"
MIN_TARGET_SIZE=$((TOTAL_USED + 1024))  # 需要1GB额外空间
info_msg "最小目标eMMC: ${MIN_TARGET_SIZE}MB (约$((MIN_TARGET_SIZE / 1024))GB)"

# 步骤2: 智能收缩eMMC分区
info_msg "步骤2: 智能收缩eMMC分区..."

# 记录原始分区信息
parted -s ${EMMC_DEV} unit s print > "${USB_DIR}/original_partition.txt"
P1_START=$(parted -s ${EMMC_DEV} unit s print | grep "^ 1" | awk '{print $2}' | sed 's/s$//')
P1_END=$(parted -s ${EMMC_DEV} unit s print | grep "^ 1" | awk '{print $3}' | sed 's/s$//')
P2_START=$(parted -s ${EMMC_DEV} unit s print | grep "^ 2" | awk '{print $2}' | sed 's/s$//')
ORIGINAL_P2_END=$(parted -s ${EMMC_DEV} unit s print | grep "^ 2" | awk '{print $3}' | sed 's/s$//')

info_msg "分区布局: P1=${P1_START}s-${P1_END}s, P2=${P2_START}s-${ORIGINAL_P2_END}s"

# 文件系统检查
e2fsck -f -y ${EMMC_DEV}p2

# 先尝试收缩到最小值
info_msg "计算文件系统最小大小..."
resize2fs -M ${EMMC_DEV}p2 2>&1 | tee /tmp/resize.log

# 获取收缩后的实际大小
FS_SIZE_BLOCKS=$(dumpe2fs -h ${EMMC_DEV}p2 2>/dev/null | grep "^Block count:" | awk '{print $3}')
BLOCK_SIZE=$(dumpe2fs -h ${EMMC_DEV}p2 2>/dev/null | grep "^Block size:" | awk '{print $3}')

# 添加默认值防止空值
[[ -z "$FS_SIZE_BLOCKS" ]] && FS_SIZE_BLOCKS=1048576  # 默认4GB worth of blocks
[[ -z "$BLOCK_SIZE" ]] && BLOCK_SIZE=4096  # 默认4K

FS_SIZE_MB=$(( (FS_SIZE_BLOCKS * BLOCK_SIZE) / 1048576 ))

info_msg "文件系统最小大小: ${FS_SIZE_MB}MB"

# 添加缓冲空间确保系统能正常运行
SHRINK_TARGET=$((FS_SIZE_MB + 512))  # 512MB缓冲
resize2fs ${EMMC_DEV}p2 ${SHRINK_TARGET}M

# 计算新的分区结束位置(扇区)
NEW_P2_END_SECTORS=$((P2_START + (SHRINK_TARGET * 1048576 / 512)))

# 重新创建收缩的分区
parted -s ${EMMC_DEV} rm 2
parted -s ${EMMC_DEV} unit s mkpart primary ext4 ${P2_START}s ${NEW_P2_END_SECTORS}s

partprobe ${EMMC_DEV}
sleep 2

info_msg "✅ eMMC分区已收缩到 ${SHRINK_TARGET}MB"

# 步骤3: 创建备份信息文件
info_msg "步骤3: 创建备份信息..."

cat > "${OUTPUT_FILE}.info" <<EOF
#!/bin/bash
# 备份信息文件 - 用于智能恢复
BACKUP_DATE="$(date)"
SOURCE_SIZE_GB=${ORIGINAL_SIZE_GB}
DATA_USED_MB=${TOTAL_USED}
MIN_TARGET_SIZE_MB=${MIN_TARGET_SIZE}

# 分区信息
P1_START=${P1_START}
P1_END=${P1_END}
P2_START=${P2_START}
P2_SHRUNK_END=${NEW_P2_END_SECTORS}
ORIGINAL_P2_END=${ORIGINAL_P2_END}
EOF

# 步骤4: 备份收缩后的eMMC
info_msg "步骤4: 备份收缩后的eMMC..."
# 确保包含完整的GPT备份表（通常在磁盘末尾）
# GPT需要首尾各34个扇区，再加一些缓冲
BACKUP_SIZE_SECTORS=$((NEW_P2_END_SECTORS + 34 + 2048))  # 34扇区GPT + 1MB额外缓冲
BACKUP_SIZE_MB=$((BACKUP_SIZE_SECTORS * 512 / 1048576))

info_msg "备份大小: ${BACKUP_SIZE_MB}MB (原始: ${ORIGINAL_SIZE_GB}GB, 可恢复到: ≥${MIN_TARGET_SIZE}MB)"

dd if=${EMMC_DEV} of="${OUTPUT_FILE}" bs=512 count=${BACKUP_SIZE_SECTORS} status=progress

# 修复GPT备份表
info_msg "修复备份镜像的GPT表..."
sgdisk -e "${OUTPUT_FILE}" 2>/dev/null || true

info_msg "✅ eMMC备份完成"

# 步骤5: 恢复eMMC原始大小
info_msg "步骤5: 恢复eMMC原始分区大小..."
parted -s ${EMMC_DEV} rm 2
parted -s ${EMMC_DEV} unit s mkpart primary ext4 ${P2_START}s ${ORIGINAL_P2_END}s

partprobe ${EMMC_DEV}
sleep 2

resize2fs ${EMMC_DEV}p2

info_msg "✅ eMMC已恢复原始大小 ${ORIGINAL_SIZE_GB}GB"

# 验证备份镜像
info_msg "步骤6: 验证备份镜像..."
LOOP_DEV=$(losetup -P -f --show "${OUTPUT_FILE}")
sleep 3

if [[ -b "${LOOP_DEV}p1" ]] && [[ -b "${LOOP_DEV}p2" ]]; then
    info_msg "✅ 备份镜像分区表完整"
    parted -s "${LOOP_DEV}" print
    
    # 检查文件系统
    e2fsck -n "${LOOP_DEV}p1" && echo "✅ Boot分区文件系统正常"
    e2fsck -n "${LOOP_DEV}p2" && echo "✅ Root分区文件系统正常"
else
    warn_msg "❌ 备份镜像分区表有问题"
fi

losetup -d "${LOOP_DEV}"

FINAL_SIZE=$(du -h "${OUTPUT_FILE}" | cut -f1)
FINAL_SIZE_ACTUAL=$(du --apparent-size -h "${OUTPUT_FILE}" | cut -f1)

# 创建智能恢复脚本
RESTORE_SCRIPT="${USB_DIR}/restore_adaptive.sh"
cat > "${RESTORE_SCRIPT}" <<'RESTORE_EOF'
#!/bin/bash
# 自适应恢复脚本 - 支持不同大小的eMMC

set -e
IMAGE_FILE="$1"
TARGET_DEV="/dev/mmcblk0"

[[ -f "${IMAGE_FILE}" ]] || { echo "用法: $0 <镜像文件>"; exit 1; }
[[ -f "${IMAGE_FILE}.info" ]] || { echo "缺少信息文件"; exit 1; }

source "${IMAGE_FILE}.info"

echo "恢复镜像到eMMC"
echo "目标设备: ${TARGET_DEV}"

TARGET_SIZE=$(blockdev --getsize64 ${TARGET_DEV})
TARGET_SIZE_MB=$((TARGET_SIZE / 1048576))

echo "目标eMMC: $((TARGET_SIZE / 1073741824))GB"

if [[ ${TARGET_SIZE_MB} -lt ${MIN_TARGET_SIZE_MB} ]]; then
    echo "错误: 目标eMMC太小，需要至少 ${MIN_TARGET_SIZE_MB}MB"
    exit 1
fi

echo "警告: 将清空 ${TARGET_DEV}!"
read -p "继续? (yes): " confirm
[[ "$confirm" == "yes" ]] || exit 1

echo "写入镜像..."
dd if="${IMAGE_FILE}" of=${TARGET_DEV} bs=4M status=progress
sync

partprobe ${TARGET_DEV}
sleep 3

# 扩展分区到最大
MAX_P2_END=$((TARGET_SIZE / 512 - 34))
parted -s ${TARGET_DEV} rm 2
parted -s ${TARGET_DEV} unit s mkpart primary ext4 ${P2_START}s ${MAX_P2_END}s
partprobe ${TARGET_DEV}
sleep 2

e2fsck -f -y ${TARGET_DEV}p2
resize2fs ${TARGET_DEV}p2

echo "✅ 恢复完成! 分区已扩展到最大空间"
RESTORE_EOF

chmod +x "${RESTORE_SCRIPT}"

info_msg "========================================="
info_msg "🎉 智能eMMC备份成功完成!"
info_msg "========================================="
info_msg "📁 备份文件: ${OUTPUT_FILE}"
info_msg "📏 备份大小: ${FINAL_SIZE}"
info_msg "📄 信息文件: ${OUTPUT_FILE}.info"
info_msg "🔧 恢复脚本: ${RESTORE_SCRIPT}"
info_msg ""
info_msg "源eMMC: ${ORIGINAL_SIZE_GB}GB"
info_msg "数据使用: ${TOTAL_USED}MB"
info_msg "最小目标: ${MIN_TARGET_SIZE}MB (约$((MIN_TARGET_SIZE / 1024))GB)"
info_msg ""
info_msg "🚀 恢复方法:"
info_msg "   1. 从SD卡启动目标设备"
info_msg "   2. 运行: bash ${RESTORE_SCRIPT} ${OUTPUT_FILE}"
info_msg ""
info_msg "✅ 支持恢复到任意大小eMMC (≥${MIN_TARGET_SIZE}MB)"
info_msg "✅ 自动扩展分区到最大可用空间"
info_msg "✅ 保持双分区结构不变"
info_msg "========================================="