#!/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="/home/umbrel/umbrel"
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"

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
}

running_off_sdcard() {
  if df -h | grep --silent /dev/mmcblk0
  then
    echo "true"
  else
    echo "false"
  fi
}

# 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* /sys/block/nvme0n* 2>/dev/null || true) | sed 's!.*/!!'
}

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

  if [[ "${device}" == nvme* ]]; then
    vendor=$(cat "/sys/block/${device}/device/device/vendor")
  else
    vendor=$(cat "/sys/block/${device}/device/vendor")
  fi
  model=$(cat "/sys/block/${device}/device/model")

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

# Returns the path of the first partition of a block device
get_first_partition_path () {
  device_path="${1}"
  if [[ "${device_path}" == /dev/nvme* ]]; then
    echo "${device_path}p1"
  else
    echo "${device_path}1"
  fi
}

# Check if a block device contains partition layout that looks like
# an umbrelOS install
is_device_umbrelos_install () {
  device_path="${1}"

  # Having a 7th partition with the label "data" is a good indicator
  # this is an umbrelOS install
  if blkid | grep "${device_path}.*7: " | grep --silent 'LABEL="data"'
  then
    echo "true"
  else
    echo "false"
  fi
}

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=$(get_first_partition_path "${device_path}")
  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 "Creating Umbrel data directory on external storage..."
  mkdir -p "${EXTERNAL_UMBREL_ROOT}"

  echo "Touching dotfile on external storage so we recognise this device in the future..."
  touch "${EXTERNAL_UMBREL_ROOT}"/.umbrel
}

# 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 () {
  echo "Running external storage mount script..."
  check_root
  check_dependencies sed wipefs parted mount sync umount

  if [[ "$(running_off_sdcard)" == "false" ]]
  then
    echo "This script should only run when umbrelOS boots from an SD card, exiting..."
    exit
  fi

  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"
      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"
    exit 1
  fi

  # Check if device is running with uas driver, if so balcklist it and reboot
  echo "Checking if we need to blacklist UAS"
  blacklist_uas_output=$(umbreld blacklist-uas || true)
  # Check output includes text "mount-script-halt" which means we are rebooting and should not mount
  if [[ "${blacklist_uas_output}" == *"mount-script-halt"* ]]; then
    echo "UAS was blacklisted and device is rebooting, exiting"
    return
  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=$(get_first_partition_path "${block_device_path}")
  block_device_model=$(get_block_device_model $block_device)
  echo "Found device \"${block_device_model}\""

  echo "Checking if the device contains an umbrelOS install..."
  if [[ $(is_device_umbrelos_install "${block_device_path}") == "true" ]]
  then
    echo "Yes, it looks like this device is an umbrelOS install not a data drive, refusing to wipe it and exiting..."
    exit 1
  fi
  echo "No, it's not"

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

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

    echo "Checking filesystem for corruption..."
    # Swallow non-zero exit because successful recovery returns 1 and
    # we still want to attempt a best effort boot on a failed recovery.
    fsck.ext4 -y "${partition_path}" || true

    echo "Mounting partition..."
    mount_partition "${partition_path}" || {
      echo "Error mounting partition"
      exit 1
    }

    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..."
  mkdir -p "${UMBREL_ROOT}" # Create the directory if it doesn't exist
  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\|/dev/nvme'

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

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

  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" | grep --quiet "overlay"

  echo "Mount script completed successfully!"
}

main
