#!/usr/bin/env bash
# vim:et:ts=2:sts=2:sw=2
# shellcheck disable=SC2218

# Get absolute repository root (especially when symlinked)
ROOT="$(realpath "$(dirname "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")")")"

# Load accompanying bash library
source "$ROOT/lib/lib.bash"

DISK="$1"
HOSTNAME="${2:-"$(uuidgen | cut -d- -f2)"}"
VERSION="${3:-"$("$ROOT/bin/tl-comp-latest" desktop)"}"
WIFI_NAME="${4:-"${WIFI_NAME:-"TechLit Africa"}"}"
WIFI_PASSWORD="${5:-"${WIFI_PASSWORD:-"empowerwatoto"}"}"
ROLE="${ROLE:-"client"}"

if [ -z "$DISK" ] || [ -z "$HOSTNAME" ] || [ -n "$HELP" ] || [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
  say "
  USAGE:
  $(g "${BASH_SOURCE##*/}")  $(c /dev/sdx)  [$(c hostname)] [ $(g 18.x.x) ]

  DESCRIPTION:
    $(g "${BASH_SOURCE##*/}") creates partitions and installs components
    for a TechLit system on $(c /dev/sdx)

    If $(c hostname)=server, server components will be installed

  ENV OPTIONS:
    $(c WIFI_NAME)=TechLit Africa
    $(c WIFI_PASSWORD)=empowerwatoto
    $(c BASTION_PORT)=
    $(c SIM_PIN)=

    $(c PREFIX)=\$(mktemp -d)
    $(c STAGE)=PREFIX
  "
  exit 1
fi

# Exit now if run as root user
assert-is-user

say "This script will need root access via $(c sudo)"
sudo echo thank you

PREFIX="${PREFIX:-"$(mktemp -d)"}" # Mountpoint to use
STAGE="${STAGE:-"$PREFIX"}"        # Prefix alternative for unpacking tarball

# Exit immediately if disk is mounted
if mount | grep -q "$DISK"; then
  err "$DISK is mounted, aborting...
If you are sure it's not the USB drive, unmount it and try again
"
  exit 1
fi

while true; do
  err "Warning!!!

Install TechLit System in $c $DISK?

All the data on $DISK will be lost!!!
$(y "Continue? Y/N")"
  TERM=xterm read -n1 -t 15 -s -r x
  case "${x:-y}" in
    y | Y) break ;;
    n | N) exit 1 ;;
  esac
done

while true; do
  err "
You will lose the$c Public storage$r and your$c config $r!!!"
  say "$(y "Proceed anyway? Y/N")"
  TERM=xterm read -n1 -t 15 -s -r x
  case "${x:-y}" in
    y | Y) break ;;
    n | N) exit 1 ;;
  esac
done

# Check for sufficient disk space before proceeding (for minis)
DISK_SIZE_B=$(sudo blockdev --getsize64 "$DISK")

if [[ "$DISK_SIZE_B" -lt 34359738368 ]]; then # 32GiB
  SRV_AVAIL_GB=$(df --output=avail -B1G /srv | tail -n 1)
  if [[ "$SRV_AVAIL_GB" -lt 15 ]]; then
    DISK_SIZE_GB=$((DISK_SIZE_B / 1024 / 1024 / 1024))

    if [ -d /srv/stage ]; then
      read -p "Target disk ($DISK) is ${DISK_SIZE_GB}GiB. /srv has only ${SRV_AVAIL_GB}GiB available. Delete /srv/stage to free up space? (y/N) " -n 1 -r
      echo
      if [[ $REPLY =~ ^[Yy]$ ]]; then
        run sudo rm -rf /srv/stage
        SRV_AVAIL_GB=$(df --output=avail -B1G /srv | tail -n 1)
      fi
    fi

    if [[ "$SRV_AVAIL_GB" -lt 15 ]]; then
      err "Still not enough space. /srv has only ${SRV_AVAIL_GB}GiB available. Please free up more space and try again."
      exit 1
    fi
  fi
fi

if ! is-t2-chip; then
  say "Zeroing disk"
  run sudo dd if="/dev/zero" of="$DISK" status=progress bs=1M count=10

  say "Partitioning and labelling"
  {
    run sudo parted "$DISK" -s mktable gpt || true
    run sudo parted "$DISK" -s mkpart primary 0 1M || true
    run sudo parted "$DISK" -s mkpart primary fat32 1M 101M || true
    run sudo parted "$DISK" -s mkpart primary 101M 102M || true
    run sudo parted "$DISK" -s mkpart primary btrfs 102M 100% || true
    run sudo parted "$DISK" -s set 1 bios_grub on || true
    run sudo parted "$DISK" -s set 2 esp on || true
  }
else
  say "T2 chip detected keeping the EFI partition"
fi

say "Waiting for kernel to recognize new partitions"
{
  part_name=$(echo "$EFI_PART" | sed -e 's;/dev/;;g')
  while ! lsblk | grep -q "$part_name"; do
    run sudo partprobe "$DISK"
    run sleep 0.5
  done
}

# Set the root partition
set-partitions "$DISK"

say "Creating filesystems"
{
  if ! is-t2-chip; then
    run sudo mkfs.fat -F32 "$EFI_PART"
  fi
  run sudo mkfs.btrfs -f "$ROOT_PART"
}

say "Mounting root filesystem (and un-mounting on errors or exit)"
{
  cleanup() {
    say "Un-mounting root filesystem"
    run sudo umount -qR "$PREFIX"
  }

  run sudo mount "$ROOT_PART" "$PREFIX"
}

say "Creating BTRFS subvolumes"
{
  run sudo btrfs subvol create "$PREFIX/@root"
  run sudo btrfs subvol create "$PREFIX/@srv"
  run sudo btrfs subvol create "$PREFIX/@swap"
  run sudo btrfs subvol create "$PREFIX/@guest"

  run sudo mkdir -p "$PREFIX/@root/home"
  run sudo mkdir -p "$PREFIX/@root/boot/efi"
  run sudo mkdir -p "$PREFIX/@root/etc"
  run sudo mkdir -p "$PREFIX/@root/srv"
  run sudo mkdir -p "$PREFIX/@root/var/btrfs"
}

say "Creating swapfile"
{
  # if memory installed is below 2G, then use 2G swap, otherwise use Mem + 2G swap
  # e.g. 1G=2G, 2G=2G, 3G=5G, 4G=6G, etc.
  SWAP="${SWAP:-"$(free -t --giga | grep Mem | awk '{print($2>2)?$2+2:2}')g"}"
  run sudo btrfs filesystem mkswapfile --size "$SWAP" "$PREFIX/@swap/swapfile"
}

say "Installing blob components"
COMPS=("snapshots" "recovery" "admin" "type" "curriculum")
if is-mini "$PREFIX"; then
  COMPS+=("desktop-mini")
else
  COMPS+=("desktop")
fi

for comp in "${COMPS[@]}"; do
  run sudo mkdir -p "$PREFIX/@srv/$comp"
done

say "Installing null components"
for comp in public secure; do
  run sudo mkdir -p "$PREFIX/@srv/$comp"
done

say "Configuring public storage"
{
  run sudo chmod 755 "$PREFIX/@srv/public"

  for group in grade-{4..8} teachers; do
    run sudo mkdir -p -m 1777 "$PREFIX/@srv/public/$group"
  done
}

say "Configuring secure storage"
{
  run sudo chmod 700 "$PREFIX/@srv/secure"
  run sudo mkdir -p -m 700 "$PREFIX/@srv/secure/ssh"
  run sudo mkdir -p -m 700 "$PREFIX/@srv/secure/net/system-connections"
  run sudo mkdir -p -m 700 "$PREFIX/@srv/secure/net/lib"
  run sudo touch "$PREFIX/@srv/secure/known_hosts"
  run sudo touch "$PREFIX/@srv/secure/bash_history"
}

say "Creating SSH identity"
{
  dir="$PREFIX/@srv/secure/ssh"
  for algo in ecdsa ed25519 rsa; do
    key="$dir/ssh_host_${algo}_key"
    if [ -f "$key" ]; then
      say " >> Destroying existing keys: $(c "$key{,.pub}")"
      run sudo rm "$key"{,.pub}
    fi
    run sudo ssh-keygen -f "$key" -N "" -t "$algo" -C "admin@$HOSTNAME"
  done
}

say "Creating persistent system config"
{
  [ "$ROLE" = "server" ] && role="server" || role="client"

  run root-write "$PREFIX/@srv/secure/config.json" "$(jq -n "{
    role: \"$role\",
    hostname: \"$HOSTNAME\",
    wifi_name: \"$WIFI_NAME\",
    wifi_password: \"$WIFI_PASSWORD\",
    gpu_hacks: \"defaults\",
    wifi_hacks: \"defaults\",
    system_hacks: \"defaults\",
  }")"
}

say "Generating fstab file"
{
  root_uuid="$(get_uuid "$ROOT_PART")"
  efi_uuid="$(get_uuid "$EFI_PART")"

  run root-write "$PREFIX/@root/etc/fstab" "UUID=$root_uuid           /              btrfs rw,relatime,subvol=@root          0 1
UUID=$root_uuid           /srv           btrfs rw,relatime,subvol=@srv           0 2
UUID=$root_uuid           /var/btrfs     btrfs rw,relatime                       0 2
UUID=$efi_uuid            /boot/efi      vfat  rw,relatime                       0 2
/var/btrfs/@swap/swapfile none           swap  sw                                0 0
tmpfs                     /tmp           tmpfs rw,mode=1777,nosuid,nodev,noatime 0 0"
}

say "Changing ownership"
{
  run sudo chown -R "$ADMIN_UID:$ADMIN_GID" "$PREFIX/@srv"
  run sudo chown -R "$ADMIN_UID:$ADMIN_GID" "$PREFIX/@root/srv"

  run sudo chown -R "$GUEST_UID:$GUEST_GID" "$PREFIX/@guest"
}

cleanup
cleanup() { :; }

# Handoff to tl-hw-update-desktop
if [ -n "$HALT" ]; then
  say "Disk prepared! Halting..."
else
  run "$ROOT/bin/tl-hw-update-desktop" "$DISK" "$VERSION"
fi
