#!/bin/bash

# SPDX-License-Identifier: GPL-2.0
# Copyright (C) 2019-present Team LibreELEC (https://libreelec.tv)

# ** !!! WARNING !!! WARNING !!! WARNING !!! WARNING !!! WARNING !!! **
# This script is not supported on generic box devices. If you choose to
# modify it to run on other hardware and you brick your device you will
# receive no support to restore your device. You have been warned!

do_backup(){
  do_umount
  echo "info: compressing ${EMMC} to ${2}, this will take time"
  dd if="${EMMC}" | gzip > of="${2}"
  echo "info: compressing completed"
}

do_detect(){
  # used on LibreComputer LePotato/LaFrite where pre-formatted modules
  # must be attached after boot (LibreComputer supports this) else the
  # device always boots to the pre-installed OS on the module.
  echo "info: rebinding d0074000.mmc"
  echo -n d0074000.mmc > /sys/bus/platform/drivers/meson-gx-mmc/unbind
  echo -n d0074000.mmc > /sys/bus/platform/drivers/meson-gx-mmc/bind
  sleep 1
  parted -s "${EMMC}" print
}

do_info(){
  BOOTINFO=$(dd if="${BOOT}" count=2048 2>/dev/null | strings | grep -e '^U-Boot ' | awk '{print $2}')
  EMMCINFO=$(dd if="${EMMC}" count=2048 2>/dev/null | strings | grep -e '^U-Boot ' | awk '{print $2}')
  echo ""
  if [ -n "${BOOTINFO}" ]; then
    echo "info: boot device is $BOOT, U-boot version is ${BOOTINFO}"
  else
    echo "info: boot device is $BOOT"
  fi
  if [ -n "${EMMCINFO}" ]; then
    echo "info: emmc device is $EMMC, U-boot version is ${EMMCINFO}"
  else
    echo "info: emmc device is $EMMC"
  fi
  echo ""
  parted -s "${EMMC}" print
  echo ""
}

do_labels(){
  if [ -n "${2}" ] && [ -n "${3}" ]; then
    BOOTLABEL=$(echo "${2}" | awk 'BEGIN { getline; print toupper($0) }')
    DISKLABEL=$(echo "${3}" | awk 'BEGIN { getline; print toupper($0) }')
  else
    BOOTLABEL="BOOT"
    DISKLABEL="DISK"
  fi
  echo "info: using boot=LABEL=${BOOTLABEL} and disk=LABEL=${DISKLABEL}"
  do_umount
  echo "info: changing label values on partitions"
  fatlabel "${EMMC}p1" "${BOOTLABEL}"
  e2label "${EMMC}p2" "${DISKLABEL}"
  sleep 1
  echo "info: remounting ${EMMC}p1 to /var/media/${BOOTLABEL}"
  mkdir -p "/var/media/${BOOTLABEL}"
  mount "${EMMC}p1" "/var/media/${BOOTLABEL}"
  echo "info: remounting ${EMMC}p2 to /var/media/${DISKLABEL}"
  mkdir -p "/var/media/${DISKLABEL}"
  mount "${EMMC}p2" "/var/media/${DISKLABEL}"
  sleep 1
  echo "info: changing label values in extlinux.conf"
  sed -i "s/boot=LABEL=\w*/boot=LABEL=${BOOTLABEL}/g" "/var/media/${BOOTLABEL}/extlinux/extlinux.conf"
  sed -i "s/disk=LABEL=\w*/disk=LABEL=${DISKLABEL}/g" "/var/media/${BOOTLABEL}/extlinux/extlinux.conf"
  sleep 1
  echo ""
  blkid | grep "${EMMC}p"
  echo ""
  cat "/var/media/${BOOTLABEL}/extlinux/extlinux.conf"
  echo ""
}

do_resize(){
  if [ -z "${DISKLABEL}" ]; then
    DISKLABEL=$(blkid | grep "${EMMC}"p2 | awk -F'"' '{print $2}')
  fi
  do_umount
  echo "info: resizing partition ${EMMC}p2 to 100%"
  parted -s -m "${EMMC}" resizepart 2 100% || parted -s -m "${EMMC}" resizepart 2 yes 100%
  sleep 5
  do_umount
  echo "info: checking filesystem"
  e2fsck -f "${EMMC}p2"
  sleep 1
  echo "info: resizing filesystem"
  resize2fs "${EMMC}p2"
  sync
  sleep 1
  echo "info: remounting ${EMMC}p2 to /var/media/${DISKLABEL}"
  mkdir -p "/var/media/${DISKLABEL}"
  mount "${EMMC}p2" "/var/media/${DISKLABEL}"
  sleep 1
  if [ -f "/var/media/${DISKLABEL}/.please_resize_me" ]; then
    rm "/var/media/${DISKLABEL}/.please_resize_me"
  fi
  parted -s "${EMMC}" print
}

do_storage(){
  if [ -n "${2}" ]; then
    DISKLABEL=$(echo "${2}" | awk 'BEGIN { getline; print toupper($0) }')
  else
    DISKLABEL="DISK"
  fi
  do_umount
  echo "info: converting emmc for /storage use"
  parted -s "${EMMC}" mklabel gpt
  parted -s "${EMMC}" -a min unit s mkpart EMMC_STORAGE ext4 34 100%
  sleep 1
  mkfs.ext4 "${EMMC}p1"
  sleep 1
  e2label "${EMMC}p1" EMMC_STORAGE
  echo "info: changing label values in extlinux.conf"
  mount -o remount,rw /flash
  sleep 1
  sed -i "s/disk=LABEL=\w*/disk=LABEL=${DISKLABEL}/g" "/flash/extlinux/extlinux.conf"
  mount -o remount,ro /flash
  sleep 1
  echo ""
  blkid | grep "${EMMC}p"
  echo ""
  cat "/flash/extlinux/extlinux.conf"
  echo ""
}

do_umount(){
  for mount in $(grep "${EMMC}" /proc/mounts | awk '{print $1}'); do
    echo "info: unmounting $mount"
    umount -f "$mount"
    sleep 2
  done
}

do_writeprotect(){
  echo  "info: disabling emmc write protection"
  echo 0 > "/sys/block/${EMMC}boot0/force_ro"
  echo 0 > "/sys/block/${EMMC}boot1/force_ro"
}

do_write(){
  do_umount

  case $(dtname) in
    radxa,zero*)
      do_writeprotect
      ;;
  esac

  if [ -e "${2}" ]; then
    case "${2}" in
      *box.img.gz|*box.img)
        echo "error: ${2} is not a bootable image, aborting!"
        exit 1
        ;;
      *.img.gz)
        echo "info: writing ${2} to ${EMMC}"
        gunzip -c "${2}" | dd of="${EMMC}" bs=1M
        ;;
      *.img)
        echo "info: writing ${2} to ${EMMC}"
        dd if="${2}" of="${EMMC}" bs=1M
        ;;
      *)
        echo "error: ${2} is not a valid image file!"
        exit 1
    esac
  else
    echo "error: ${2} not found!"
    exit 1
  fi
}

do_zero(){
  do_umount
  echo "info: zeroing ${EMMC}, be patient, this will take time"
  dd if="/dev/zero" of="${EMMC}" bs=1M
  echo "info: zeroing complete"
}

do_help(){
  echo ""
  echo "usage: emmctool (w)rite <filename>    : write <filename>.img/.img.gz to the eMMC module"
  echo "                (b)backup <filename>  : dump the emmc partition to <filename>.img.gz file"
  echo "                (d)etect              : detect an eMMC module attached after boot"
  echo "                (i)nfo                : show info about the eMMC module"
  echo "                (l)abel <boot> <disk> : change eMMC disk labels to <bootlabel> <disklabel>"
  echo "                (r)esize              : resize the storage partition to 100%"
  echo "                (s)storage            : convert emmc for use as /storage (boot from sdcard)"
  echo "                (z)ero                : zero (erase/wipe) the eMMC module"
  echo "                (h)elp                : displays this help message"
  echo ""
}

BOOT=$(grep /flash /proc/mounts | awk '{print $1}' | sed 's/p[012]//g')
EMMC=$(find /dev -name "mmcblk*rpmb" | sed 's/rpmb//g' | head -n 1)

if [ -z "${EMMC}" ]; then
  echo "error: no emmc module detected!"
  exit 1
fi

case $(dtname) in
  azw*|bananapi*|friendlyarm*|hardkernel*|khadas*|libretech*|radxa*|wetek*)
    if [ "${BOOT}" = "${EMMC}" ]; then
      do_info
      echo "Your device is booted from the eMMC module!"
      echo ""
      exit 1
    fi
    ;;
  *)
    echo "Your device is not supported!"
    echo ""
    exit 1
    ;;
esac

case "${1}" in
  backup)
    do_backup "$@"
    ;;
  detect|d)
    do_detect
    ;;
  info|i)
    do_info
    ;;
  labels|l)
    do_labels "$@"
    ;;
  resize|r)
    do_resize
    ;;
  storage|s)
    do_storage "$@"
    ;;
  write|w)
    do_write "$@"
    do_resize
    do_labels
    ;;
  zero|z)
    do_zero
    ;;
  *)
    do_info
    do_help
    ;;
esac

exit
