#!/usr/bin/env bash
set -o errexit
set -o nounset
set -o pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

run_base=0
run_shell=0
run_cockpit=0
run_docker=0
run_libvirt=0
run_nvidia=0
run_tuning=0
run_security=0
enable_ufw=0
nvidia_driver_version=""

usage() {
  cat <<'EOF'
Usage: sudo ./install.sh [OPTIONS]

Day-0 bootstrap automation for Ubuntu 24.04 or newer.

Profiles:
  --base                         Install baseline operational packages
  --shell                        Install the root shell profile
  --cockpit                      Install and enable Cockpit
  --docker                       Install and configure Docker
  --libvirt                      Install and enable libvirt/KVM
  --nvidia-tools                 Install NVIDIA diagnostic tools only
  --install-nvidia-driver VERSION
                                 Install diagnostic tools and the explicit driver
  --tuning                       Install journald and sysctl tuning
  --security                     Install fail2ban and UFW; do not enable UFW
  --enable-ufw                   Run security profile and explicitly enable UFW
  --all                          Run every profile without enabling UFW or
                                 installing an NVIDIA driver
  -h, --help                     Show this help

Examples:
  sudo ./install.sh --base --shell
  sudo ./install.sh --all
  sudo ./install.sh --all --enable-ufw
  sudo ./install.sh --nvidia-tools --install-nvidia-driver 550
EOF
}

require_supported_ubuntu() {
  if [[ ! -r /etc/os-release ]]; then
    printf 'CRITICAL: /etc/os-release is unavailable; refusing system changes\n' >&2
    exit 2
  fi

  # shellcheck disable=SC1091
  source /etc/os-release
  if [[ "${ID:-}" != "ubuntu" ]]; then
    printf 'CRITICAL: this toolkit supports Ubuntu only; detected %s\n' "${ID:-unknown}" >&2
    exit 2
  fi
  if ! dpkg --compare-versions "${VERSION_ID:-0}" ge "24.04"; then
    printf 'CRITICAL: Ubuntu 24.04 or newer is required; detected %s\n' \
      "${VERSION_ID:-unknown}" >&2
    exit 2
  fi
}

if (($# == 0)); then
  usage
  exit 0
fi

while (($# > 0)); do
  case "$1" in
    --base)
      run_base=1
      ;;
    --shell)
      run_shell=1
      ;;
    --cockpit)
      run_cockpit=1
      ;;
    --docker)
      run_docker=1
      ;;
    --libvirt)
      run_libvirt=1
      ;;
    --nvidia-tools)
      run_nvidia=1
      ;;
    --install-nvidia-driver)
      if (($# < 2)); then
        printf 'CRITICAL: --install-nvidia-driver requires a VERSION\n' >&2
        exit 2
      fi
      nvidia_driver_version="$2"
      if [[ ! "$nvidia_driver_version" =~ ^[0-9]+$ ]]; then
        printf 'CRITICAL: NVIDIA driver VERSION must contain digits only\n' >&2
        exit 2
      fi
      run_nvidia=1
      shift
      ;;
    --tuning)
      run_tuning=1
      ;;
    --security)
      run_security=1
      ;;
    --enable-ufw)
      enable_ufw=1
      run_security=1
      ;;
    --all)
      run_base=1
      run_shell=1
      run_cockpit=1
      run_docker=1
      run_libvirt=1
      run_nvidia=1
      run_tuning=1
      run_security=1
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      printf 'CRITICAL: unknown option: %s\n\n' "$1" >&2
      usage >&2
      exit 2
      ;;
  esac
  shift
done

if ((EUID != 0)); then
  printf 'CRITICAL: install.sh must run as root for selected profiles\n' >&2
  exit 2
fi

for required_command in bash dpkg; do
  if ! command -v "$required_command" >/dev/null 2>&1; then
    printf 'CRITICAL: required command is missing: %s\n' "$required_command" >&2
    exit 2
  fi
done

require_supported_ubuntu

printf 'INFO: running read-only preflight\n'
"$SCRIPT_DIR/scripts/00-preflight.sh"

((run_base == 0)) || "$SCRIPT_DIR/scripts/01-base-packages.sh"
((run_shell == 0)) || "$SCRIPT_DIR/scripts/02-shell-profile.sh"
((run_cockpit == 0)) || "$SCRIPT_DIR/scripts/03-cockpit.sh"
((run_docker == 0)) || "$SCRIPT_DIR/scripts/04-docker.sh"
((run_libvirt == 0)) || "$SCRIPT_DIR/scripts/05-libvirt.sh"

if ((run_nvidia == 1)); then
  if [[ -n "$nvidia_driver_version" ]]; then
    "$SCRIPT_DIR/scripts/06-nvidia-tools.sh" --install-driver "$nvidia_driver_version"
  else
    "$SCRIPT_DIR/scripts/06-nvidia-tools.sh"
  fi
fi

((run_tuning == 0)) || "$SCRIPT_DIR/scripts/07-tuning.sh"

if ((run_security == 1)); then
  if ((enable_ufw == 1)); then
    "$SCRIPT_DIR/scripts/08-security-baseline.sh" --enable-ufw
  else
    "$SCRIPT_DIR/scripts/08-security-baseline.sh"
  fi
fi

printf '\nINFO: running post-install checks\n'
"$SCRIPT_DIR/scripts/99-postcheck.sh"
printf '\nOK: selected Linux setup profiles completed\n'