#!/usr/bin/env bash

set -euo pipefail

# This script will:
# - Look for external storage devices
# - Check if they contain an Umbrel install
# - If yes
# - - Mount it
# - If no
# - - Format it
# - - Mount it
# - - Install Umbrel on it
# - Bind mount the external installation on top of the local installation

UMBREL_ROOT="$(readlink -f $(dirname "${BASH_SOURCE[0]}")/../../..)"
MOUNT_POINT="/mnt/data"
EXTERNAL_UMBREL_ROOT="${MOUNT_POINT}/umbrel"
DOCKER_DIR="/var/lib/docker"
EXTERNAL_DOCKER_DIR="${MOUNT_POINT}/docker"
SWAP_DIR="/swap"
SWAP_FILE="${SWAP_DIR}/swapfile"

set_status="/status-server/set-status"

check_root () {
  if [[ $UID != 0 ]]; then
    echo "This script must be run as root"
    exit 1
  fi
}

check_dependencies () {
  for cmd in "$@"; do
    if ! command -v $cmd >/dev/null 2>&1; then
      echo "This script requires \"${cmd}\" to be installed"
      exit 1
    fi
  done
}

# Returns a list of block device paths
list_block_devices () {
  # We need to run sync here to make sure the filesystem is reflecting the
  # the latest changes in /sys/block/sd*
  sync
  # We use "2>/dev/null || true" to swallow errors if there are
  # no block devices. In that case the function just returns nothing
  # instead of an error which is what we want.
  #
  # sed 's!.*/!!' is to return the device path so we get sda
  # instead of /sys/block/sda
  (ls -d /sys/block/sd* 2>/dev/null || true) | sed 's!.*/!!'
}

# Returns the vendor and model name of a block device
get_block_device_model () {
  device="${1}"
  vendor=$(cat "/sys/block/${device}/device/vendor")
  model=$(cat "/sys/block/${device}/device/model")

  # We echo in a subshell without quotes to strip surrounding whitespace
  echo "$(echo $vendor) $(echo $model)"
}

# By default Linux uses the UAS driver for most devices. This causes major
# stability problems on the Raspberry Pi, not due to issues with UAS, but due
# to devices running in UAS mode using much more power. The Pi can't reliably
# provide enough power to the USB port and the entire system experiences
# extreme instability. By blacklisting all devices from the UAS driver on boot
# we fall back to the mass-storage driver, which results in decreased
# performance, but lower power usage, and much better system stability.
blacklist_uas () {
  usb_quirks=$(lsusb | awk '{print $6":u"}' | tr '\n' ',' | sed 's/,$//')
  echo -n "${usb_quirks}" > /sys/module/usb_storage/parameters/quirks

  echo "Rebinding USB drivers..."
  for i in /sys/bus/pci/drivers/[uoex]hci_hcd/*:*; do
    [[ -e "$i" ]] || continue;
    echo "${i##*/}" > "${i%/*}/unbind"
    echo "${i##*/}" > "${i%/*}/bind"
  done
}

is_partition_ext4 () {
  partition_path="${1}"
  # We need to run sync here to make sure the filesystem is reflecting the
  # the latest changes in /dev/*
  sync
  blkid -o value -s TYPE "${partition_path}" | grep --quiet '^ext4$'
}

# Wipes a block device and reformats it with a single EXT4 partition
format_block_device () {
  device="${1}"
  device_path="/dev/${device}"
  partition_path="${device_path}1"
  wipefs -a "${device_path}"
  parted --script "${device_path}" mklabel gpt
  parted --script "${device_path}" mkpart primary ext4 0% 100%
  # We need to run sync here to make sure the filesystem is reflecting the
  # the latest changes in /dev/*
  sync
  mkfs.ext4 -F -L umbrel "${partition_path}"
}

# Mounts the device given in the first argument at $MOUNT_POINT
mount_partition () {
  partition_path="${1}"
  mkdir -p "${MOUNT_POINT}"
  mount "${partition_path}" "${MOUNT_POINT}"
}

# Unmounts $MOUNT_POINT
unmount_partition () {
  umount "${MOUNT_POINT}"
}

# Formats and sets up a new device
setup_new_device () {
  block_device="${1}"
  partition_path="${2}"

  echo "Formatting device..."
  format_block_device $block_device

  echo "Mounting partition..."
  mount_partition "${partition_path}"

  echo "Copying Umbrel install to external storage..."
  mkdir -p "${EXTERNAL_UMBREL_ROOT}"
  cp  --recursive \
      --archive \
      --no-target-directory \
      "${UMBREL_ROOT}" "${EXTERNAL_UMBREL_ROOT}"
}

# Copy Docker data dir to external storage
copy_docker_to_external_storage () {
  mkdir -p "${EXTERNAL_DOCKER_DIR}"
  cp  --recursive \
      --archive \
      --no-target-directory \
      "${DOCKER_DIR}" "${EXTERNAL_DOCKER_DIR}"
}

main () {
  $set_status mount started
  echo "Running external storage mount script..."
  check_root
  check_dependencies sed wipefs parted mount sync umount

  no_of_block_devices=$(list_block_devices | wc -l)

  retry_for_block_devices=1

  while [[ $no_of_block_devices -lt 1 ]]; do

    echo "No block devices found"
    echo "Waiting for 5 seconds before checking again..."

    sleep 5

    no_of_block_devices=$(list_block_devices | wc -l)
    retry_for_block_devices=$(( $retry_for_block_devices + 1 ))

    if [[ $retry_for_block_devices -gt 20 ]]; then
      echo "No block devices found in 20 tries..."
      echo "Exiting mount script without doing anything"
      $set_status mount errored no-block-device
      exit 1
    fi

  done

  if [[ $no_of_block_devices -gt 1 ]]; then
    echo "Multiple block devices found, only one drive is supported"
    echo "Exiting mount script without doing anything"
    $set_status mount errored multiple-block-devices
    exit 1
  fi

  # At this point we know there is only one block device attached
  block_device=$(list_block_devices)
  block_device_path="/dev/${block_device}"
  partition_path="${block_device_path}1"
  block_device_model=$(get_block_device_model $block_device)
  echo "Found device \"${block_device_model}\""

  echo "Blacklisting USB device IDs against UAS driver..."
  blacklist_uas

  echo "Checking USB devices are back..."
  retry_for_usb_devices=1
  while [[ ! -e "${block_device_path}" ]]; do
    retry_for_usb_devices=$(( $retry_for_usb_devices + 1 ))
    if [[ $retry_for_usb_devices -gt 10 ]]; then
      echo "USB devices weren't registered after 10 tries..."
      echo "Exiting mount script without doing anything"
      $set_status mount errored rebinding-failed
      exit 1
    fi

    echo "Waiting for USB devices..."
    sleep 1
  done

  echo "Checking if the device is ext4..."

  if is_partition_ext4 "${partition_path}" ; then
    echo "Yes, it is ext4"
    mount_partition "${partition_path}"

    echo "Checking if device contains an Umbrel install..."

    if [[ -f "${EXTERNAL_UMBREL_ROOT}"/.umbrel ]]; then
      echo "Yes, it contains an Umbrel install"
    else
      echo "No, it doesn't contain an Umbrel install"
      echo "Unmounting partition..."
      unmount_partition
      setup_new_device $block_device $partition_path
    fi

  else
    echo "No, it's not ext4"
    setup_new_device $block_device $partition_path
  fi

  if [[ ! -d "${EXTERNAL_DOCKER_DIR}" ]]; then
    echo "Copying Docker data directory to external storage..."
    copy_docker_to_external_storage
  fi

  echo "Bind mounting external storage over local Umbrel installation..."
  mount --bind "${EXTERNAL_UMBREL_ROOT}" "${UMBREL_ROOT}"

  echo "Bind mounting external storage over local Docker data dir..."
  mount --bind "${EXTERNAL_DOCKER_DIR}" "${DOCKER_DIR}"

  echo "Bind mounting external storage to ${SWAP_DIR}"
  mkdir -p "${MOUNT_POINT}/swap" "${SWAP_DIR}"
  mount --bind "${MOUNT_POINT}/swap" "${SWAP_DIR}"

  echo "Bind mounting SD card root at /sd-card..."
  [[ ! -d "/sd-root" ]] && mkdir -p "/sd-root"
  mount --bind "/" "/sd-root"

  echo "Checking Umbrel root is now on external storage..."
  sync
  sleep 1
  df -h "${UMBREL_ROOT}" | grep --quiet '/dev/sd'

  echo "Checking ${DOCKER_DIR} is now on external storage..."
  df -h "${DOCKER_DIR}" | grep --quiet '/dev/sd'

  echo "Checking ${SWAP_DIR} is now on external storage..."
  df -h "${SWAP_DIR}" | grep --quiet '/dev/sd'

  echo "Setting up swapfile"
  rm "${SWAP_FILE}" || true
  fallocate -l 4G "${SWAP_FILE}"
  chmod 600 "${SWAP_FILE}"
  mkswap "${SWAP_FILE}"
  swapon "${SWAP_FILE}"

  echo "Checking SD Card root is bind mounted at /sd-root..."
  df -h "/sd-root${UMBREL_ROOT}" | grep --quiet "/dev/root"

  echo "Starting external drive mount monitor..."
  echo
  ${UMBREL_ROOT}/scripts/umbrel-os/external-storage/monitor ${block_device} ${MOUNT_POINT} &

  echo "Mount script completed successfully!"
  $set_status mount completed
}

main