#!/usr/bin/env bash
# @file Quick Start Provision Script
# @brief Main entry point for Install Doctor that ensures Homebrew and a few dependencies are installed before cloning the repository and running Chezmoi.
# @description
#     This script ensures Homebrew is installed and then installs a few dependencies that Install Doctor relies on.
#     After setting up the minimal amount of changes required, it clones the Install Doctor repository (which you
#     can customize the location of so you can use your own fork). It then proceeds by handing things over to
#     Chezmoi which handles the dotfile application and synchronous scripts. Task is used in conjunction with
#     Chezmoi to boost the performance in some spots by introducing asynchronous features.
#
#     **Note**: `https://install.doctor/start` points to this file.
#
#     ## Dependencies
#
#     The chart below shows the dependencies we rely on to get Install Doctor going. The dependencies that are bolded
#     are mandatory. The ones that are not bolded are conditionally installed only if they are required.
#
#     | Dependency         | Description                                                                          |
#     |--------------------|--------------------------------------------------------------------------------------|
#     | **Chezmoi**        | Dotfile configuration manager (on-device provisioning)                               |
#     | **Task**           | Task runner used on-device for task parallelization and dependency management        |
#     | **ZX / Node.js**   | ZX is a Node.js abstraction that allows for better scripts                           |
#     | Gum                | Gum is a terminal UI prompt CLI (which allows sweet, interactive prompts)            |
#     | Glow               | Glow is a markdown renderer used for applying terminal-friendly styled to markdown   |
#
#     There are also a handful of system packages that are installed like `curl` and `git`. Then, during the Chezmoi provisioning
#     process, there are a handful of system packages that are installed to ensure things run smoothly. You can find more details
#     about these extra system packages by browsing through the `home/.chezmoiscripts/${DISTRO_ID}/` folder and other applicable
#     folders (e.g. `universal`).
#
#     Although Install Doctor comes with presets that install a whole gigantic amount of software, it can actually
#     be quite good at provisioning minimal server environments where you want to keep the binaries to a minimum.
#
#     ## Variables
#
#     Specify certain environment variables to customize the behavior of Install Doctor. With the right combination of
#     environment variables, this script can be run completely headlessly. This allows us to do things like test our
#     provisioning script on a wide variety of operating systems.
#
#     | Variable                  | Description                                                                       |
#     |---------------------------|-----------------------------------------------------------------------------------|
#     | `START_REPO` (or `REPO`)  | Variable to specify the Git fork to use when provisioning                         |
#     | `ANSIBLE_PROVISION_VM`    | **For Qubes**, determines the name of the VM used to provision the system         |
#     | `DEBUG_MODE` (or `DEBUG`) | Set to true to enable verbose logging                                             |
#
#     For a full list of variables you can use to customize Install Doctor, check out our [Customization](https://install.doctor/docs/customization)
#     and [Secrets](https://install.doctor/docs/customization/secrets) documentation.
#
#     ## Links
#
#     [Install Doctor homepage](https://install.doctor)
#     [Install Doctor documentation portal](https://install.doctor/docs) (includes tips, tricks, and guides on how to customize the system to your liking)

{{ include "partials" "logg" }}
# @section Environment variables and system dependencies
# @description Ensure Ubuntu / Debian run in `noninteractive` mode. Detect `START_REPO` format and determine appropriate git address,
#     otherwise use the master Install Doctor branch
setEnvironmentVariables() {
  export DEBIAN_FRONTEND=noninteractive
  export HOMEBREW_NO_ENV_HINTS=true
  if [ -z "$START_REPO" ] && [ -z "$REPO" ]; then
    export START_REPO="https://github.com/megabyte-labs/install.doctor.git"
  else
    if [ -n "$REPO" ] && [ -z "$START_REPO" ]; then
      export START_REPO="$REPO"
    fi
    if [[ "$START_REPO" == *"/"* ]]; then
      # Either full git address or GitHubUser/RepoName
      if [[ "$START_REPO" == *":"* ]] || [[ "$START_REPO" == *"//"* ]]; then
        export START_REPO="$START_REPO"
      else
        export START_REPO="https://github.com/${START_REPO}.git"
      fi
    else
      export START_REPO="https://github.com/$START_REPO/install.doctor.git"
    fi
  fi
}

{{ include "partials" "basic-deps" }}
{{ include "partials" "homebrew" }}
{{ include "partials" "reboot" }}
# @description Load default settings if it is in a CI setting
setCIEnvironmentVariables() {
  if [ -n "$CI" ]; then
    logg info "Automatically setting environment variables since the CI environment variable is defined"
    logg info "Setting HOST to $HOST" && export HOST="$HOST"
    logg info "Setting NO_RESTART to true" && export NO_RESTART=true
    logg info "Setting HEADLESS_INSTALL to true " && export HEADLESS_INSTALL=true
    logg info "Setting SOFTWARE_GROUP to Full-Desktop" && export SOFTWARE_GROUP="Full-Desktop"
    logg info "Setting FULL_NAME to Brian Zalewski" && export FULL_NAME="Brian Zalewski"
    logg info "Setting PRIMARY_EMAIL to brian@megabyte.space" && export PRIMARY_EMAIL="brian@megabyte.space"
    logg info "Setting PUBLIC_SERVICES_DOMAIN to lab.megabyte.space" && export PUBLIC_SERVICES_DOMAIN="lab.megabyte.space"
    logg info "Setting RESTRICTED_ENVIRONMENT to false" && export RESTRICTED_ENVIRONMENT=false
    logg info "Setting WORK_ENVIRONMENT to false" && export WORK_ENVIRONMENT=false
  fi
}

# @description Disconnect from WARP, if connected
ensureWarpDisconnected() {
  if command -v warp-cli > /dev/null; then
    if warp-cli status | grep 'Connected' > /dev/null; then
      logg info "Disconnecting from WARP" && warp-cli disconnect && logg success "Disconnected WARP to prevent conflicts"
    fi
  fi
}

# @description Notify user that they can press CTRL+C to prevent `/etc/sudoers` from being modified (which is currently required for headless installs on some systems).
#     Additionally, this function will add the current user to `/etc/sudoers` so that headless automation is possible.
setupPasswordlessSudo() {
  sudo -n true || SUDO_EXIT_CODE=$?
  logg info 'Your user will temporarily be granted passwordless sudo for the duration of the script'
  if [ -n "$SUDO_EXIT_CODE" ]; then
    logg info 'Press CTRL+C to bypass this prompt to either enter your password when needed or perform a non-privileged installation'
    logg info 'Note: Non-privileged installations are not yet supported'
  fi
  if [ -n "$SUDO_PASSWORD" ]; then
    printf '%s\n' "$SUDO_PASSWORD" | sudo -p "" -S echo "$(whoami) ALL=(ALL:ALL) NOPASSWD: ALL # TEMPORARY FOR INSTALL DOCTOR" | sudo tee -a /etc/sudoers > /dev/null
  else
    echo "$(whoami) ALL=(ALL:ALL) NOPASSWD: ALL # TEMPORARY FOR INSTALL DOCTOR" | sudo tee -a /etc/sudoers > /dev/null
  fi
}

# @section Qubes dom0
# @description Ensure sys-whonix is configured (for Qubes dom0)
ensureSysWhonix() {
  CONFIG_WIZARD_COUNT=0
  function configureWizard() {
    if xwininfo -root -tree | grep "Anon Connection Wizard"; then
      WINDOW_ID="$(xwininfo -root -tree | grep "Anon Connection Wizard" | sed 's/^ *\([^ ]*\) .*/\1/')"
      xdotool windowactivate "$WINDOW_ID" && sleep 1 && xdotool key 'Enter' && sleep 1 && xdotool key 'Tab Tab Enter' && sleep 24 && xdotool windowactivate "$WINDOW_ID" && sleep 1 && xdotool key 'Enter' && sleep 300
      qvm-shutdown --wait sys-whonix
      sleep 3
      qvm-start sys-whonix
      if xwininfo -root -tree | grep "systemcheck | Whonix" > /dev/null; then
        WINDOW_ID_SYS_CHECK="$(xwininfo -root -tree | grep "systemcheck | Whonix" | sed 's/^ *\([^ ]*\) .*/\1/')"
        if xdotool windowactivate "$WINDOW_ID_SYS_CHECK"; then
          sleep 1
          xdotool key 'Enter'
        fi
      fi
    else
      sleep 3
      CONFIG_WIZARD_COUNT=$((CONFIG_WIZARD_COUNT + 1))
      if [[ "$CONFIG_WIZARD_COUNT" == '4' ]]; then
        echo "The sys-whonix anon-connection-wizard utility did not open."
      else
        echo "Checking for anon-connection-wizard again.."
        configureWizard
      fi
    fi
  }
}

# @description Ensure dom0 is updated
ensureDom0Updated() {
  if [ ! -f /root/dom0-updated ]; then
    sudo qubesctl --show-output state.sls update.qubes-dom0
    sudo qubes-dom0-update --clean -y
    touch /root/dom0-updated
  fi
}

# @description Ensure sys-whonix is running
ensureSysWhonixRunning() {
  if ! qvm-check --running sys-whonix; then
    qvm-start sys-whonix --skip-if-running
    configureWizard > /dev/null
  fi
}

# @description Ensure TemplateVMs are updated
ensureTemplateVMsUpdated() {
  if [ ! -f /root/templatevms-updated ]; then
    # timeout of 10 minutes is added here because the whonix-gw VM does not like to get updated
    # with this method. Anyone know how to fix this?
    sudo timeout 600 qubesctl --show-output --skip-dom0 --templates state.sls update.qubes-vm &> /dev/null || true
    while read -r RESTART_VM; do
      qvm-shutdown --wait "$RESTART_VM"
    done< <(qvm-ls --all --no-spinner --fields=name,state | grep Running | grep -v sys-net | grep -v sys-firewall | grep -v sys-whonix | grep -v dom0 | awk '{print $1}')
    sudo touch /root/templatevms-updated
  fi
}

# @description Ensure provisioning VM can run commands on any VM
ensureProvisioningVMPermissions() {
  echo "/bin/bash" | sudo tee /etc/qubes-rpc/qubes.VMShell
  sudo chmod 755 /etc/qubes-rpc/qubes.VMShell
  echo "${ANSIBLE_PROVISION_VM:=provision}"' dom0 allow' | sudo tee /etc/qubes-rpc/policy/qubes.VMShell
  echo "$ANSIBLE_PROVISION_VM"' $anyvm allow' | sudo tee -a /etc/qubes-rpc/policy/qubes.VMShell
  sudo chown "$(whoami):$(whoami)" /etc/qubes-rpc/policy/qubes.VMShell
  sudo chmod 644 /etc/qubes-rpc/policy/qubes.VMShell
}

# @description Create provisioning VM and initialize the provisioning process from there
createAndInitProvisionVM() {
  qvm-create --label red --template debian-11 "$ANSIBLE_PROVISION_VM" &> /dev/null || true
  qvm-volume extend "$ANSIBLE_PROVISION_VM:private" "40G"
  if [ -f ~/.vaultpass ]; then
    qvm-run "$ANSIBLE_PROVISION_VM" 'rm -f ~/QubesIncoming/dom0/.vaultpass'
    qvm-copy-to-vm "$ANSIBLE_PROVISION_VM" ~/.vaultpass
    qvm-run "$ANSIBLE_PROVISION_VM" 'cp ~/QubesIncoming/dom0/.vaultpass ~/.vaultpass'
  fi
}

# @description Restart the provisioning process with the same script but from the provisioning VM
runStartScriptInProvisionVM() {
  qvm-run --pass-io "$ANSIBLE_PROVISION_VM" 'curl -sSL https://install.doctor/start > ~/start.sh && bash ~/start.sh'
}

# @description Perform Qubes dom0 specific logic like updating system packages, setting up the Tor VM, updating TemplateVMs, and
#     beginning the provisioning process using Ansible and an AppVM used to handle the provisioning process
handleQubesDom0() {
  if command -v qubesctl > /dev/null; then
    ensureSysWhonix
    ensureDom0Updated
    ensureSysWhonixRunning
    ensureTemplateVMsUpdated
    ensureProvisioningVMPermissions
    createAndInitProvisionVM
    runStartScriptInProvisionVM
    exit 0
  fi
}

# @section Homebrew dependencies
# @description Helper function used by [[ensureHomebrewDeps]] to ensure a Homebrew package is installed after
#     first checking if it is already available on the system.
installBrewPackage() {
  if ! command -v "$1" > /dev/null; then
    logg 'Installing '"$1"''
    brew install --quiet "$1"
  fi
}

# @description Installs various dependencies using Homebrew.
#
#     1. Ensures Chezmoi, Node.js, and ZX are installed.
#     2. Installs Glow and Gum if the `HEADLESS_INSTALL` environment variable is not set.
#     3. If the system is macOS, then also install `gsed` and `coreutils`.
ensureHomebrewDeps() {
  ### TUI experience
  if [ -z "$HEADLESS_INSTALL" ]; then
    installBrewPackage "gum"
    installBrewPackage "glow"
  fi

  ### Base dependencies
  installBrewPackage "chezmoi"
  installBrewPackage "node"
  installBrewPackage "zx"

  ### macOS
  if [ -d /Applications ] && [ -d /System ]; then
    installBrewPackage "expect"
    installBrewPackage "gsed"
    if ! command -v gtimeout > /dev/null; then
      brew install --quiet coreutils
    fi
  fi
}

# @section Chezmoi
# @description Ensure the `${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi` directory is cloned and up-to-date using the previously
#     set `START_REPO` as the source repository.
cloneChezmoiSourceRepo() {
  if ! git config --get --global http.postBuffer > /dev/null; then
    logg info 'Setting git http.postBuffer value high for large source repository' && git config --global http.postBuffer 524288000
  fi
  if [ -d "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi/.git" ]; then
    logg info "Changing directory to ${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi" && cd "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi"
    logg info "Pulling the latest changes in ${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi" && git pull origin master
  else
    logg info "Ensuring ${XDG_DATA_HOME:-$HOME/.local/share} is a folder" && mkdir -p "${XDG_DATA_HOME:-$HOME/.local/share}"
    logg info "Cloning ${START_REPO} to ${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi" && git clone "${START_REPO}" "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi"
  fi
}

# @description Guide the user through the initial setup by showing TUI introduction and accepting input through various prompts.
#
#     1. Show `chezmoi-intro.md` with `glow`
#     2. Prompt for the software group if the `SOFTWARE_GROUP` variable is not defined
#     3. Run `chezmoi init` when the Chezmoi configuration is missing (i.e. `${XDG_CONFIG_HOME:-$HOME/.config}/chezmoi/chezmoi.yaml`)
initChezmoiAndPrompt() {
  ### Show `chezmoi-intro.md` with `glow`
  if command -v glow > /dev/null; then
    glow "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi/docs/terminal/chezmoi-intro.md"
  fi

  ### Prompt for the software group if the `SOFTWARE_GROUP` variable is not defined
  if command -v gum > /dev/null; then
    if [ -z "$SOFTWARE_GROUP" ]; then
      logg prompt 'Select the software group you would like to install. If your environment is a macOS, Windows, or environment with the DISPLAY environment variable then desktop software will be installed too. The software groups are in the '"${XDG_CONFIG_HOME:-$HOME/.config}/chezmoi/chezmoi.yaml"' file.'
      SOFTWARE_GROUP="$(gum choose "Basic" "Server" "Standard" "Full")"
      export SOFTWARE_GROUP
    fi
  else
    logg error 'Woops! Gum needs to be installed for the guided installation. Try running brew install gum' && exit 1
  fi

  if [ ! -f "${XDG_CONFIG_HOME:-$HOME/.config}/chezmoi/chezmoi.yaml" ]; then
    ### Run `chezmoi init` when the Chezmoi configuration is missing
    logg info 'Running chezmoi init since the '"${XDG_CONFIG_HOME:-$HOME/.config}/chezmoi/chezmoi.yaml"' is not present'
    chezmoi init
  fi
}

# @description Run `chezmoi apply` and enable verbose mode if the `DEBUG_MODE` or `DEBUG` environment variable is set to true
configureDebugMode() {
  if [ -n "$DEBUG_MODE" ] || [ -n "$DEBUG" ]; then
    logg info "Either DEBUG_MODE or DEBUG environment variables were set so Chezmoi will be run in debug mode"
    export DEBUG_MODIFIER="-vvvvv"
  fi
}

# @description Save the log of the provision process to `$HOME/.local/var/log/install.doctor/install.doctor.$(date +%s).log` and add the Chezmoi
#     `--force` flag if the `HEADLESS_INSTALL` variable is set to `true`.
runChezmoi() {
  mkdir -p "$HOME/.local/var/log/install.doctor"
  LOG_FILE="$HOME/.local/var/log/install.doctor/install.doctor.$(date +%s).log"
  if [ "$HEADLESS_INSTALL" = 'true' ]; then
    logg info 'Running chezmoi apply forcefully'
    if command -v unbuffer > /dev/null; then
      if command -v caffeinate > /dev/null; then
        caffeinate unbuffer -p chezmoi apply $DEBUG_MODIFIER -k --force 2>&1 | tee "$LOG_FILE"
      else
        unbuffer -p chezmoi apply $DEBUG_MODIFIER -k --force 2>&1 | tee "$LOG_FILE"
      fi
    else
      if command -v caffeinate > /dev/null; then
        caffeinate chezmoi apply $DEBUG_MODIFIER -k --force 2>&1 | tee "$LOG_FILE"
      else
        chezmoi apply $DEBUG_MODIFIER -k --force 2>&1 | tee "$LOG_FILE"
      fi
    fi
  else
    logg info 'Running chezmoi apply'
    if command -v unbuffer > /dev/null; then
      if command -v caffeinate > /dev/null; then
        caffeinate unbuffer -p chezmoi apply $DEBUG_MODIFIER -k 2>&1 | tee "$LOG_FILE"
      else
        unbuffer -p chezmoi apply $DEBUG_MODIFIER -k 2>&1 | tee "$LOG_FILE"
      fi    
    else
      if command -v caffeinate > /dev/null; then
        caffeinate chezmoi apply $DEBUG_MODIFIER -k 2>&1 | tee "$LOG_FILE"
      else
        chezmoi apply $DEBUG_MODIFIER -k 2>&1 | tee "$LOG_FILE"
      fi    
    fi
  fi
}

# @section Post-provision logic
# @description Ensure temporary passwordless sudo privileges are removed from `/etc/sudoers`
removePasswordlessSudo() {
  if command -v gsed > /dev/null; then
    sudo gsed -i '/# TEMPORARY FOR INSTALL DOCTOR/d' /etc/sudoers || logg warn 'Failed to remove passwordless sudo from the /etc/sudoers file'
  else
    sudo sed -i '/# TEMPORARY FOR INSTALL DOCTOR/d' /etc/sudoers || logg warn 'Failed to remove passwordless sudo from the /etc/sudoers file'
  fi
}

# @description Render the `docs/terminal/post-install.md` file to the terminal at the end of the provisioning process
postProvision() {
  logg success 'Provisioning complete!'
  if command -v glow > /dev/null && [ -f "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi/docs/terminal/post-install.md" ]; then
    glow "${XDG_DATA_HOME:-$HOME/.local/share}/chezmoi/docs/terminal/post-install.md"
  fi
}


# @section Execution order
# @description The `provisionLogic` function is used to define the order of the script. All of the functions it relies on are defined
#     above.
provisionLogic() {
  logg info "Setting environment variables" && setEnvironmentVariables
  logg info "Handling CI variables" && setCIEnvironmentVariables
  logg info "Ensuring WARP is disconnected" && ensureWarpDisconnected
  logg info "Applying passwordless sudo" && setupPasswordlessSudo
  logg info "Ensuring system Homebrew dependencies are installed" && ensureBasicDeps
  logg info "Ensuring Homebrew is available" && ensureHomebrew
  logg info "Installing Homebrew packages" && ensureHomebrewDeps
  logg info "Handling Qubes dom0 logic (if applicable)" && handleQubesDom0
  logg info "Cloning / updating source repository" && cloneChezmoiSourceRepo
  logg info "Handling pre-provision logic" && initChezmoiAndPrompt
  logg info "Handling debug mode if DEBUG or DEBUG_MODE are defined" && configureDebugMode
  logg info "Running the Chezmoi provisioning" && runChezmoi
  logg info "Ensuring temporary passwordless sudo is removed" && removePasswordlessSudo
  logg info "Handling post-provision logic" && postProvision
  logg info "Determing whether or not reboot" && handleRequiredReboot
}
provisionLogic
