#!/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

while true; do
  err "WARNING: This will wipe the disk "$DISK" (Including everything in the public storage)"
  say "$(y Proceed?) $(c '[Y/n]')"
  read -n1 -s -r x
  case "${x:-y}" in
    y | Y) break ;;
    n | N) exit 1 ;;
  esac
done

# 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

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"
for comp in snapshots desktop recovery admin; 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
