#!/bin/bash
function checkVariables {
  if [[ -z $DISK ]]; then
    printf "System drive (DISK) is not set\n"
    exit 99
  fi

  if [[ -z $HOSTNAME ]]; then
    printf "Hostname (HOSTNAME) is not set\n"
    exit 99
  fi

  if [[ -z $USERNAME ]]; then
    printf "Username (USERNAME) is not set\n"
    exit 99
  fi

  if [[ -z $PASSWORD ]]; then
    printf "Password (PASSWORD) is not set\n"
    exit 99
  fi
}

function preChrootStep {
  if [[ $(whoami) != "root" ]]; then
    printf "This script has to be run as root"
    exit 1
  fi

  if ! ls /sys/firmware/efi/efivars > /dev/null 2>&1; then
    printf "This script can only be run when booted in UEFI mode\n"
    exit 1
  fi

  if ! ping -c 1 archlinux.org > /dev/null 2>&1; then
    printf "Unable to ping archlinux.org, check internet connectivity\n"
    exit 1
  fi

  printf "Arch installation\n"
  read -p "System drive: " DISK
  read -p "Hostname: " HOSTNAME
  read -p "Username: " USERNAME
  PASSWORD=$(openssl passwd -6) || exit $?
  printf "Generate ssh key\n"
  ssh-keygen -f /tmp/id_ed25519 -t ed25519 -C "${USERNAME}@${HOSTNAME}" -q || exit $?
  SSH_PUB=$(cat /tmp/id_ed25519.pub)
  SSH_PRIV=$(cat /tmp/id_ed25519)

  LOGFILE="archInstall.$(date +%Y%m%d-%H%M%S).log"

  printf "Installing arch on ${DISK}\n" | tee -a $LOGFILE

  BLOCK_INFO=$(lsblk -n --output PATH,TYPE)
  printf "$BLOCK_INFO" | grep -e " disk$" | grep -e "^${DISK} " > /dev/null 2>&1
  if (( $? > 0 )); then
    printf "${DISK} does not seem to be a disk\n" | tee -a $LOGFILE
    exit 2
  fi

  if (( $(printf "$BLOCK_INFO" | grep -e " part$" | wc -l) > 0 )); then
    printf "${DISK} already has several partitions\n" | tee -a $LOGFILE
    read -p "Do you really wish to continue and erase all partitions? [N/y] " CONTINUE
    if [[ ! $CONTINUE =~ ^[yY]$ ]]; then
      printf "Exiting due to user input\n" | tee -a $LOGFILE
      exit 0
    fi
  fi

  printf "Loading keymap\n" | tee -a $LOGFILE
  loadkeys sv-latin1 || exit $?

  printf "Creating partitions\n" | tee -a $LOGFILE
  (
    echo o      # New partition table
    echo n      # New partition
    echo p      # Primary partition
    echo 1      # Partition number
    echo        # First sector (use default)
    echo +300M  # Last sector
    echo n      # New partition
    echo p      # Primary partition
    echo 2      # Partition number
    echo        # First sector (use default)
    echo        # Last sector (use default)
    echo w      # Write changes
  ) | fdisk -W always ${DISK} >> $LOGFILE 2>&1 || exit $?

  BOOTPART="${DISK}1"
  BTRFSPART="${DISK}2"

  printf "Creating filesystems\n" | tee -a $LOGFILE
  mkfs.fat ${BOOTPART} >> $LOGFILE 2>&1 || exit $?
  mkfs.btrfs ${BTRFSPART} >> $LOGFILE 2>&1 || exit $?

  printf "Mounting btrfs partition\n" | tee -a $LOGFILE
  mount ${BTRFSPART} /mnt >> $LOGFILE 2>&1 || exit $?

  printf "Creating btrfs subvolumes\n" | tee -a $LOGFILE
  btrfs subvolume create /mnt/@root >> $LOGFILE 2>&1 || exit $?
  btrfs subvolume create /mnt/@home >> $LOGFILE 2>&1 || exit $?
  btrfs subvolume create /mnt/@log >> $LOGFILE 2>&1 || exit $?
  btrfs subvolume create /mnt/@swap >> $LOGFILE 2>&1 || exit $?

  printf "Unmounting btrfs partition\n" | tee -a $LOGFILE
  umount /mnt >> $LOGFILE 2>&1

  printf "Mounting root\n" | tee -a $LOGFILE
  mount -o defaults,relatime,compress=zstd,subvol=@root ${BTRFSPART} /mnt >> $LOGFILE 2>&1 || exit $?

  printf "Creating mount directories\n" | tee -a $LOGFILE
  mkdir -p /mnt/boot/efi /mnt/home /mnt/var/log /mnt/swap >> $LOGFILE 2>&1 || exit $?

  printf "Mounting volumes\n" | tee -a $LOGFILE
  mount ${BOOTPART} /mnt/boot/efi >> $LOGFILE 2>&1 || exit $?
  mount ${BTRFSPART} -o defaults,relatime,compress=zstd,subvol=@home /mnt/home >> $LOGFILE 2>&1 || exit $?
  mount ${BTRFSPART} -o defaults,relatime,compress=zstd,subvol=@log /mnt/var/log >> $LOGFILE 2>&1 || exit $?
  mount ${BTRFSPART} -o defaults,relatime,compress=zstd,subvol=@swap /mnt/swap >> $LOGFILE 2>&1 || exit $?

  printf "Creating swap file\n" | tee -a $LOGFILE
  btrfs filesystem mkswapfile --size 8G /mnt/swap/swapfile >> $LOGFILE 2>&1 || exit $?
  swapon /mnt/swap/swapfile

  PACSTRAPPKGS="base base-devel linux linux-firmware btrfs-progs grub efibootmgr networkmanager sudo sed git ansible"

  printf "Checking CPU manufacturer\n" | tee -a $LOGFILE
  CPU=$(lscpu | grep "^Vendor ID:" | awk '{ print $3 }')

  if [[ ! -z $CPU ]]; then
    if [[ "$CPU" == "GenuineIntel" ]]; then
      PACSTRAPPKGS="${PACSTRAPPKGS} intel-ucode"
    elif [[ "$CPU" == "AuthencticAMD" ]]; then
      PACSTRAPPKGS="${PACSTRAPPKGS} amd-ucode"
    fi
  fi

  printf "Installing base system\n" | tee -a $LOGFILE
  pacstrap -K /mnt ${PACSTRAPPKGS} >> $LOGFILE 2>&1 || exit $?

  printf "Generate fstab\n" | tee -a $LOGFILE
  genfstab -U /mnt >> /mnt/etc/fstab || exit $?
}

function chrootStep {
  checkVariables

  printf "Setting up time\n"
  ln -sf /usr/share/zoneinfo/Europe/Stockholm /etc/localtime || exit $?
  hwclock --systohc || exit $?

  printf "Setting up locale\n"
  sed -i -e 's/^#\(en_US.UTF-8\)/\1/' /etc/locale.gen || exit $?
  sed -i -e 's/^#\(sv_SE.UTF-8\)/\1/' /etc/locale.gen || exit $?
  locale-gen >&2 || exit $?

  echo "LANG=en_US.UTF-8" > /etc/locale.conf
  echo "LC_TIME=sv_SE.UTF-8" >> /etc/locale.conf
  echo "KEYMAP=sv-latin1" > /etc/vconsole.conf

  printf "Setting hostname to $HOSTNAME\n"
  echo "$HOSTNAME" > /etc/hostname

  printf "Add wheel to sudoers\n"
  echo "%wheel ALL=(ALL) ALL" > /etc/sudoers.d/wheel
  sed -i -e 's/^#\(%wheel ALL=(ALL) ALL\)/\1/' /etc/sudoers || exit $?

  printf "Creating user\n" | tee -a $LOGFILE
  useradd -m $USERNAME -G wheel >&2 || exit $?
  echo "${USERNAME}:${PASSWORD}" | chpasswd -e >&2 || exit $?

  printf "Add user ssh key" | tee -a $LOGFILE
  (umask 066; mkdir /home/${USERNAME}/.ssh)
  (umask 066; echo "${SSH_PRIV}" > /home/${USERNAME}/.ssh/id_ed25519)
  (umask 022; echo "${SSH_PUB}" > /home/${USERNAME}/.ssh/id_ed25519.pub)
  chown -R ${USERNAME}:${USERNAME} /home/${USERNAME}/.ssh

  printf "Setting temporary root password\n"
  echo "root:root" | chpasswd >&2 || exit $?

  printf "Starting and enabling NetworkManager\n"
  systemctl enable NetworkManager >&2 || exit $?
  systemctl start NetworkManager >&2 || exit $?

  printf "Installing GRUB\n"
  grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=GRUB >&2 || exit $?
  grub-mkconfig -o /boot/grub/grub.cfg >&2 || exit $?

  printf "Exiting chroot\n"
  exit
}

function postChrootStep {
  printf "Cleanup\n" | tee -a $LOGFILE
  cp $LOGFILE /mnt/$LOGFILE
  cd /
  swapoff /mnt/swap/swapfile
  umount -R /mnt
}

if [[ $1 == "chroot" ]]; then
  chrootStep
  exit
fi

# Run preChroot operations
preChrootStep

# Run chroot operations
printf "Adding install script to mount directory\n" | tee -a $LOGFILE
cp $0 /mnt/

printf "Chrooting arch\n" | tee -a $LOGFILE
DISK="${DISK}" \
  HOSTNAME="${HOSTNAME}" \
  USERNAME="${USERNAME}" \
  PASSWORD="${PASSWORD}" \
  SSH_PUB="${SSH_PUB}" \
  SSH_PRIV="${SSH_PRIV}" \
  arch-chroot /mnt /$(basename $0) chroot 2>> $LOGFILE | tee -a $LOGFILE
if (( ${PIPESTATUS[0]} > 0 )); then
  exit ${PIPESTATUS[0]}
fi

# Run postchroot operations
postChrootStep

printf "Base installation completed, logfile at $LOGFILE.\nReboot and proceed with goodies.\nRemember to change/deactivate root login\n" | tee -a $LOGFILE