#!/usr/bin/env bash
#
## WizardKit: Live Linux Build Tool

set -o errexit
set -o errtrace
set -o nounset
set -o pipefail

# Prep
DATE="$(date +%Y-%m-%d)"
DATETIME="$(date +%Y-%m-%d_%H%M)"
ROOT_DIR="$(realpath $(dirname "$0")/..)"
BUILD_DIR="$ROOT_DIR/setup/BUILD"
LOG_DIR="$BUILD_DIR/logs"
OUT_DIR="$ROOT_DIR/setup/OUT_LINUX"
PROFILE_DIR="$BUILD_DIR/archiso-profile"
REPO_DIR="$BUILD_DIR/repo"
SKEL_DIR="$PROFILE_DIR/airootfs/etc/skel"
TEMP_DIR="$BUILD_DIR/temp"
WORK_DIR="$(mktemp -dt archiso-tmp.XXXXXXXXXX)"
ISO_DIR="${WORK_DIR}/iso"
MIRRORLIST_SOURCE='https://archlinux.org/mirrorlist/?country=US&protocol=http&protocol=https&ip_version=4&use_mirror_status=on'
if command -v nano >/dev/null 2>&1; then
  EDITOR=nano
elif command -v vim >/dev/null 2>&1; then
  EDITOR=vim
else
  EDITOR=vi
fi
if [ ! -z ${SUDO_USER+x} ]; then
  REAL_USER="$SUDO_USER"
fi

function ask() {
  while :; do
    read -p "$1 [Y/N] " -r answer
    if echo "$answer" | grep -Eiq '^(y|yes|sure)$'; then
      return 0
    elif echo "$answer" | grep -Eiq '^(n|no|nope)$'; then
      return 1
    fi
  done
}

function cleanup() {
  for d in "$TEMP_DIR" "$PROFILE_DIR"; do
    if [[ -d "$d" ]]; then
      if ask "Remove: ${d}?"; then
        rm -Rf "$d"
      fi
    fi
  done
  mkdir -p "$BUILD_DIR" 2>/dev/null
  mkdir -p "$LOG_DIR" 2>/dev/null
  mkdir -p "$OUT_DIR" 2>/dev/null
  mkdir -p "$TEMP_DIR" 2>/dev/null
}

function fix_kit_permissions() {
  # GitHub zip archives don't preserve the correct permissions
  for d in images scripts setup; do
    find "$ROOT_DIR/$d" -type d -exec chmod 755 "{}" \;
  done
}

function load_settings() {
  dos2unix "$ROOT_DIR/scripts/wk/cfg/main.py"
  while read line; do
    if echo "$line" | grep -Eq "^\w+='"; then
      line="$(echo "$line" | sed -r 's/[\r\n]+//')"
      eval "$line"
    fi
  done < "$ROOT_DIR/scripts/wk/cfg/main.py"
}

function copy_live_env() {
  echo "Copying Archlinux files..."
  rsync -aI "$ROOT_DIR/setup/linux/profile/" "$PROFILE_DIR/"
  mkdir -p "$PROFILE_DIR/airootfs/usr/local/bin"
  rsync -aI "$ROOT_DIR/scripts/" "$PROFILE_DIR/airootfs/usr/local/bin/"

  # Pre-compile Python scripts
  unset PYTHONPYCACHEPREFIX
  python -m compileall "$PROFILE_DIR/airootfs/usr/local/bin/"

  # Update profiledef.sh to set proper permissions for executable files
  for _file in $(find "$PROFILE_DIR/airootfs" -executable -type f | sed "s%$PROFILE_DIR/airootfs%%" | sort); do
    sed -i "\$i\  [\"$_file\"]=\"0:0:755\"" "$PROFILE_DIR/profiledef.sh"
  done
}

function run_elevated() {
  prog="$1"
  shift
  if command -v sudo >/dev/null 2>&1; then
    if ! sudo "$prog" $*; then
      echo "ERROR: Failed to run '$prog'"
      if ask "Retry?"; then
        sudo "$prog" $*
      fi
    fi
  else
    echo -n "Root "
    if ! su -c "export REAL_USER=$USER && '$prog' $*"; then
      echo "ERROR: Failed to run '$prog'"
      if ask "Retry?"; then
        su -c "export REAL_USER=$USER && '$prog' $*"
      fi
    fi
  fi
}

function update_live_env() {
  hostname="$(echo "${KIT_NAME_SHORT}-linux" | tr "[:upper:]" "[:lower:]")"
  username="tech"
  label="${KIT_NAME_SHORT}_LINUX"

  # MOTD
  sed -i -r "s/KIT_NAME_SHORT/$KIT_NAME_SHORT/" "$PROFILE_DIR/profiledef.sh"
  sed -i -r "s/KIT_NAME_FULL/$KIT_NAME_SHORT/" "$PROFILE_DIR/profiledef.sh"
  sed -i -r "s/SUPPORT_URL/$KIT_NAME_SHORT/" "$PROFILE_DIR/profiledef.sh"

  # Boot config (legacy)
  mkdir -p "$TEMP_DIR" 2>/dev/null
  git clone --depth=1 https://github.com/ipxe/wimboot "$TEMP_DIR/wimboot"
  rsync -aI "$TEMP_DIR/wimboot"/{LICENSE.txt,README.md,wimboot} "$PROFILE_DIR/syslinux/wimboot/"
  cp "$ROOT_DIR/images/Pxelinux.png" "$PROFILE_DIR/syslinux/pxelinux.png"
  cp "$ROOT_DIR/images/Syslinux.png" "$PROFILE_DIR/syslinux/syslinux.png"
  sed -i -r "s/__+/$KIT_NAME_FULL/" "$PROFILE_DIR/syslinux/syslinux.cfg"

  # Boot config (UEFI)
  curl -Lo "$TEMP_DIR/refind.zip" "https://sourceforge.net/projects/refind/files/latest/download"
  7z x -aoa "$TEMP_DIR/refind.zip" -o"$TEMP_DIR/refind"
  cp "$ROOT_DIR/images/rEFInd.png" "$PROFILE_DIR/EFI/boot/rEFInd.png"
  cp "$TEMP_DIR/refind"/refind*/"refind/refind_x64.efi" "$PROFILE_DIR/EFI/boot/bootx64.efi"
  rsync -aI "$TEMP_DIR/refind"/refind*/refind/drivers_x64/ "$PROFILE_DIR/EFI/boot/drivers_x64/"
  rsync -aI "$TEMP_DIR/refind"/refind*/refind/icons/ "$PROFILE_DIR/EFI/boot/icons/"
  sed -i "s/%ARCHISO_LABEL%/${label}/" "$PROFILE_DIR/EFI/boot/refind.conf"

  # Memtest86+ (Open Source)
  mkdir -p "$PROFILE_DIR/EFI/memtest86+"
  mkdir -p "$TEMP_DIR/memtest86+"
  curl -Lo "$TEMP_DIR/memtest86+/memtest86-binaries.zip" "https://memtest.org/download/v6.10/mt86plus_6.10.binaries.zip"
  7z e "$TEMP_DIR/memtest86+/memtest86-binaries.zip" -o"$TEMP_DIR/memtest86+" "memtest64.efi"
  mv "$TEMP_DIR/memtest86+/memtest64.efi" "$PROFILE_DIR/EFI/memtest86+/bootx64.efi"

  # Memtest86 (Passmark)
  mkdir -p "$PROFILE_DIR/EFI/memtest86/Benchmark"
  mkdir -p "$TEMP_DIR/memtest86"
  curl -Lo "$TEMP_DIR/memtest86/memtest86-usb.zip" "https://www.memtest86.com/downloads/memtest86-usb.zip"
  7z e -aoa "$TEMP_DIR/memtest86/memtest86-usb.zip" -o"$TEMP_DIR/memtest86" "memtest86-usb.img"
  7z e -aoa "$TEMP_DIR/memtest86/memtest86-usb.img" -o"$TEMP_DIR/memtest86" "MemTest86.img"
  7z x -aoa "$TEMP_DIR/memtest86/MemTest86.img" -o"$TEMP_DIR/memtest86"
  rm "$TEMP_DIR/memtest86/EFI/BOOT/BOOTIA32.efi"
  mv "$TEMP_DIR/memtest86/EFI/BOOT/BOOTX64.efi" "$PROFILE_DIR/EFI/memtest86/bootx64.efi"
  mv "$TEMP_DIR/memtest86/EFI/BOOT"/* "$PROFILE_DIR/EFI/memtest86"/
  mv "$TEMP_DIR/memtest86/help"/* "$PROFILE_DIR/EFI/memtest86"/
  mv "$TEMP_DIR/memtest86/license.rtf" "$PROFILE_DIR/EFI/memtest86"/

  # Hostname
  echo "$hostname" > "$PROFILE_DIR/airootfs/etc/hostname"
  echo "127.0.1.1  $hostname.localdomain  $hostname" >> "$PROFILE_DIR/airootfs/etc/hosts"

  # Live packages
  cp "$ROOT_DIR/setup/linux/packages/base" "$PROFILE_DIR/packages.x86_64"
  echo "[custom]" >> "$PROFILE_DIR/pacman.conf"
  echo "SigLevel = Optional TrustAll" >> "$PROFILE_DIR/pacman.conf"
  echo "Server = file://$REPO_DIR" >> "$PROFILE_DIR/pacman.conf"

  # Mirrors
  mkdir -p "$PROFILE_DIR/airootfs/etc/pacman.d"
  curl -Lo "$PROFILE_DIR/airootfs/etc/pacman.d/mirrorlist" "$MIRRORLIST_SOURCE"
  sed -i 's/#Server/Server/g' "$PROFILE_DIR/airootfs/etc/pacman.d/mirrorlist"

  # MOTD
  sed -i -r "s/_+/$KIT_NAME_FULL Linux Environment/" "$PROFILE_DIR/airootfs/etc/motd"

  # Network
  ln -s "/run/systemd/resolve/stub-resolv.conf" "$PROFILE_DIR/airootfs/etc/resolv.conf"

  # Oh My ZSH
  git clone --depth=1 https://github.com/robbyrussell/oh-my-zsh.git "$SKEL_DIR/.oh-my-zsh"
  rm -Rf "$SKEL_DIR/.oh-my-zsh/.git"
  curl -o "$SKEL_DIR/.oh-my-zsh/themes/lean.zsh-theme" https://raw.githubusercontent.com/miekg/lean/master/prompt_lean_setup

  # Openbox theme
  git clone --depth=1 https://github.com/addy-dclxvi/Openbox-Theme-Collections.git "$TEMP_DIR/ob-themes"
  mkdir -p "$PROFILE_DIR/airootfs/usr/share/themes"
  cp -a "$TEMP_DIR/ob-themes/Triste-Orange" "$PROFILE_DIR/airootfs/usr/share/themes/"

  # Rofi
  ## Probably don't need the exact commit but it'll be fine
  mkdir -p "$PROFILE_DIR/airootfs/usr/share/fonts/"
  curl -Lo \
    "$PROFILE_DIR/airootfs/usr/share/fonts/Fantasque-Sans-Mono-Nerd-Font.ttf" \
    "https://github.com/adi1090x/rofi/raw/9c4093c665326bb08d6affc7e16d18d8f25c4452/fonts/Fantasque-Sans-Mono-Nerd-Font.ttf"
  curl -Lo \
    "$PROFILE_DIR/airootfs/usr/share/fonts/Feather.ttf" \
    "https://github.com/adi1090x/rofi/raw/9c4093c665326bb08d6affc7e16d18d8f25c4452/fonts/Feather.ttf"

  # SSH
  mkdir -p "$SKEL_DIR/.ssh"
  ssh-keygen -b 4096 -C "$username@$hostname" -N "" -f "$SKEL_DIR/.ssh/id_rsa"
  if ! grep -qv "^#" "$ROOT_DIR/setup/linux/authorized_keys"; then
    echo "WARNING: No authorized SSH keys found." 1>&2
  fi
  cp "$ROOT_DIR/setup/linux/authorized_keys" "$SKEL_DIR/.ssh/authorized_keys"

  # Root user
  echo "root:$(echo "$ROOT_PASSWORD" | openssl passwd -6 -stdin):14871::::::" >> "$PROFILE_DIR/airootfs/etc/shadow"

  # Tech user
  echo "tech:$(echo "$TECH_PASSWORD" | openssl passwd -6 -stdin):14871::::::" >> "$PROFILE_DIR/airootfs/etc/shadow"

  # Timezone
  ln -sf "/usr/share/zoneinfo/$LINUX_TIME_ZONE" "$PROFILE_DIR/airootfs/etc/localtime"

  # VNC password
  mkdir "$SKEL_DIR/.vnc"
  echo "$TECH_PASSWORD" | vncpasswd -f > "$SKEL_DIR/.vnc/passwd"

  # Wallpaper
  mkdir -p "$PROFILE_DIR/airootfs/usr/share/wallpaper"
  cp "$ROOT_DIR/images/Linux.png" "$PROFILE_DIR/airootfs/usr/share/wallpaper/burned.in"

  # udevil
  mkdir -p "$PROFILE_DIR/airootfs/media"

  # WiFi
  IFS_BAK="${IFS}"
  IFS=$'\n'
  mkdir -p "$PROFILE_DIR/airootfs/var/lib/iwd"
  for line in $(<"$ROOT_DIR/setup/linux/known_networks"); do
    if [[ "${line:0:1}" == "#" ]]; then
      # Skip comments
      continue
    fi
    w_name="${line%%:::*}"
    w_pass="${line##*:::}"
    w_pass="$(wpa_passphrase "${w_name}" "${w_pass}" \
      | grep -E 'psk=[0-9a-z]+' \
      | sed -r 's/\s+psk=//')"
    echo "[Security]" > "$PROFILE_DIR/airootfs/var/lib/iwd/${w_name}.psk"
    echo "PreSharedKey=${w_pass}" >> "$PROFILE_DIR/airootfs/var/lib/iwd/${w_name}.psk"
  done
  IFS="${IFS_BAK}"
}

function update_repo() {
  if [[ "$EUID" -eq 0 ]]; then
    echo "This section not meant to be run as root."
    echo "Aborted."
    exit 1
  fi
  echo "Updating custom repo..."

  # Archive current files
  if [[ -d "$REPO_DIR" ]]; then
    mkdir -p "$BUILD_DIR/Archive" 2>/dev/null
    archive="$BUILD_DIR/Archive/$(date "+%Y-%m-%d_%H%M%S")"
    mv -bv "$REPO_DIR" "$archive"
  fi
  sleep 1s

  # Prep directories
  mkdir -p "$REPO_DIR" 2>/dev/null
  mkdir -p "$TEMP_DIR/repo" 2>/dev/null
  pushd "$TEMP_DIR/repo" >/dev/null

  # Prep custom repo packages
  while read -r p; do
    echo "Prepping: $p"
    curl -LsfO https://aur.archlinux.org/cgit/aur.git/snapshot/$p.tar.gz
    tar xf $p.tar.gz
    if [[ "$p" == "hfsprogs" ]]; then
      sed -i 's!http://cavan.codon.org.uk/\~mjg59/diskdev_cmds!https://sources.voidlinux.org/hfsprogs-540.1.linux3!' "$TEMP_DIR/repo/hfsprogs/PKGBUILD"
    fi
  done < "$ROOT_DIR/setup/linux/packages/aur"
  cp -an "$ROOT_DIR/setup/linux/packages/custom"/* "$TEMP_DIR/repo"/
  rm ./*.tar.gz

  # Build custom repo packages
  for p in *; do
    echo "Building: $p"
    pushd $p >/dev/null
    makepkg -d
    popd >/dev/null
    mv -n $p/*zst "$REPO_DIR"/
  done
  popd >/dev/null

  # Build custom repo database
  pushd "$REPO_DIR" >/dev/null
  repo-add custom.db.tar.gz *zst
  popd >/dev/null
}

function install_deps() {
  echo "Installing dependencies..."
  packages=
  while read -r line; do
    if ! echo "$line" | grep -Fq "#"; then
      packages="$packages $line"
    fi
  done < "$ROOT_DIR/setup/linux/packages/dependencies"
  run_elevated pacman -Syu --needed --noconfirm $packages
}

function build_linux() {
  cleanup
  fix_kit_permissions
  install_deps
  load_settings --edit

  # Update repo if necessary
  if ! [[ -e "${REPO_DIR}/custom.db.tar.gz" ]]; then
    update_repo
  fi

  # Build live environment
  copy_live_env
  update_live_env

  # Rerun script as root to start Archiso build process
  run_elevated "$(realpath "$0")" --build-iso
}

function build_iso() {
  if [[ "$EUID" -ne 0 ]]; then
    echo "This section is meant to be run as root."
    echo "Aborted."
    exit 1
  fi

  # Removing cached (and possibly outdated) custom repo packages
  for package in $(cat "$ROOT_DIR/setup/linux/packages/aur"); do
    for p in /var/cache/pacman/pkg/*${package}*; do
      if [[ -f "${p}" ]]; then
        rm "${p}"
      fi
    done
  done

  # Build ISO
  prefix="${KIT_NAME_SHORT}-Linux"
  label="${KIT_NAME_SHORT}_LINUX"
  mkarchiso \
    -o "$OUT_DIR" \
    -w ${WORK_DIR} \
    -v "$PROFILE_DIR" \
    | tee -a "$LOG_DIR/$DATETIME.log"

  # Build better ISO
  rsync -aI "$PROFILE_DIR/EFI/" "${ISO_DIR:-safety}/EFI/"
  rsync -aI --ignore-existing "$PROFILE_DIR/syslinux/" "${ISO_DIR:-safety}/syslinux/"
  ## Sketchy bit ##
  . /usr/bin/mkarchiso -o "${OUT_DIR}" -w "${WORK_DIR}" "${PROFILE_DIR}"
  isofs_dir="${ISO_DIR}"
  image_name="${KIT_NAME_SHORT}-Linux-${DATE}-x86_64.iso"
  rm "${OUT_DIR}/${image_name}"
  _build_iso_image

  # Cleanup
  echo "Removing temp files..."
  rm "$TEMP_DIR/Linux" -Rf | tee -a "$LOG_DIR/$DATETIME.log"
  sudo umount -R ${WORK_DIR} || true
  sudo rm -rf ${WORK_DIR}

  echo "Reverting permissions..."
  chown $REAL_USER:$REAL_USER "$PROFILE_DIR" -R
  chown $REAL_USER:$REAL_USER "$OUT_DIR" -R
}

# Check input
case ${1:-} in
  -b|--build)
    build_linux
    echo Done
    ;;

  -f|--fix-perms)
    fix_kit_permissions
    echo Done
    ;;

  -i|--install-deps)
    install_deps
    echo Done
    ;;

  -o|--build-iso)
    load_settings
    build_iso
    echo Done
    ;;

  -p|--prep-live-env)
    load_settings --edit
    copy_live_env
    update_live_env
    echo Done
    ;;

  -u|--update-repo)
    update_repo
    echo Done
    ;;

  *)
    echo "Usage: $(basename "$0") [OPTIONS]"
    echo ""
    echo "Options:"
    echo "    -b --build          Perform all tasks to build the iso"
    echo "    -h --help           Show usage"
    echo ""
    echo "Advanced options:"
    echo "    -f --fix-perms      Fix folder permissions"
    echo "    -i --install-deps   Install build dependencies"
    echo "    -o --build-iso      Build ISO (using current setup)"
    echo "    -p --prep-live-env  Prep live & airootfs folders"
    echo "    -u --update-repo    Update custom pacman repo"
    ;;
esac