Bash script to setup a new Ubuntu PC with my custom setup

new_ubuntu_setup.sh

#!/usr/bin/env bash
set -euo pipefail

# =========================================================
# NJ Setup Script v2.0 (modular, flag-driven, ubuntu-ready)
# =========================================================

# -------------------------------
# Defaults (all sections opt-in)
INSTALL_ALL="no"

SSH_KEYS="no"
DIRS="no"
APT_DESKTOP="no"
DEV="no"
BASHRC="no"
VSCODE="no"
PYTORCH="no"
OLLAMA="no"
DRY_RUN="no"

GPU_TYPE="cpu"

LOG_FILE="$HOME/nj_setup_$(date +%Y%m%d_%H%M%S).log"

# -------------------------------
# ASCII Banner
print_banner() {
cat << "EOF"

 _   _      _
| \ | |    (_)
|  \| |     _   ___
| . ` |    | | / __|
| |\  |    | | \__ \
|_| \_|    |_| |___/

EOF
}

# -------------------------------
# Logging
log() {
  echo -e "\n\033[1;34m[INFO]\033[0m $1" | tee -a "$LOG_FILE"
}

warn() {
  echo -e "\033[1;33m[WARN]\033[0m $1" | tee -a "$LOG_FILE"
}

die() {
  echo -e "\n\033[1;31m[ERROR]\033[0m $1" | tee -a "$LOG_FILE"
  exit 1
}

# Run commands safely (no eval). Preserves exit status even with tee.
run() {
  if [[ "$DRY_RUN" == "yes" ]]; then
    echo "[DRY_RUN] $*" | tee -a "$LOG_FILE"
    return 0
  fi

  "$@" 2>&1 | tee -a "$LOG_FILE"
  local status=${PIPESTATUS[0]}
  if [[ $status -ne 0 ]]; then
    die "Command failed ($status): $*"
  fi
}

on_err() {
  die "Something failed at line $1. Check log: $LOG_FILE"
}
trap 'on_err $LINENO' ERR

# -------------------------------
# Help / Usage
usage() {
  cat << EOF
Usage:
  $0 [flags]

Flags (yes/no):
  --install-all  Install everything (default: no)

  --ssh          Setup SSH keys (default: no)
  --dirs         Create folders (default: no)
  --apt          Install desktop/user APT packages (default: no)
  --dev          Install dev tools + create vision venv (default: no)
  --bashrc       Add NJ bashrc section (default: no)
  --vscode       Install VS Code + extensions (default: no)
  --pytorch      Install PyTorch into vision venv (default: no)
  --ollama       Install Ollama + pull model based on GPU (default: no)
  --dry-run      Print commands only (default: no)

Examples:
  $0 --install-all yes
  $0 --install-all yes --dry-run yes
  $0 --ssh yes --dirs yes --apt yes --dev yes --bashrc yes
  $0 --dev yes --pytorch yes --ollama yes
EOF
}

need_value() {
  local flag="$1"
  local val="${2:-}"
  if [[ -z "$val" ]]; then
    die "Missing value for $flag (use yes/no)."
  fi
}

validate_yesno() {
  local name="$1" val="$2"
  if [[ "$val" != "yes" && "$val" != "no" ]]; then
    die "$name must be 'yes' or 'no' (got '$val')"
  fi
}

# -------------------------------
# OS Detection
detect_os() {
  if [[ ! -f /etc/os-release ]]; then
    die "Cannot detect OS (missing /etc/os-release)."
  fi
  if ! grep -q "Ubuntu" /etc/os-release; then
    die "This script supports Ubuntu only."
  fi

  if ! command -v lsb_release &>/dev/null; then
    run sudo apt update
    run sudo apt install -y lsb-release
  fi

  local ubuntu_version
  ubuntu_version="$(lsb_release -rs)"
  log "Detected Ubuntu $ubuntu_version"

  # Proper version compare (22.04+)
  local min_version="22.04"
  if ! printf '%s\n%s\n' "$min_version" "$ubuntu_version" | sort -V -C; then
    die "Ubuntu $min_version+ required (detected $ubuntu_version)."
  fi
}

# -------------------------------
# GPU Detection
detect_gpu() {
  if command -v nvidia-smi &>/dev/null; then
    GPU_TYPE="nvidia"
    log "NVIDIA GPU detected."
  else
    GPU_TYPE="cpu"
    log "No NVIDIA GPU detected. Using CPU mode."
  fi
}

# -------------------------------
# SSH Keys
setup_ssh_keys() {
  log "Setting up SSH keys..."

  run mkdir -p "$HOME/.ssh"
  run chmod 700 "$HOME/.ssh"

  if [[ -f "$HOME/.ssh/id_ed25519" ]]; then
    log "SSH key already exists: $HOME/.ssh/id_ed25519"
    return 0
  fi

  # No passphrase by default (change if you want)
  run ssh-keygen -t ed25519 -C "$USER@$(hostname)" -N "" -f "$HOME/.ssh/id_ed25519"
  log "SSH key generated: $HOME/.ssh/id_ed25519"
}

# -------------------------------
# Folders
create_dirs() {
  log "Creating directories..."
  local dirs=(
    "$HOME/Software"
    "$HOME/NJ"
    "$HOME/Hobby/Musicmaking/Ampero"
    "$HOME/Hobby/Programming"
    "$HOME/Book"
    "$HOME/Downloads/Programs"
    "$HOME/Downloads/Documents"
    "$HOME/Downloads/Media"
    "$HOME/NJ_Github"
    "$HOME/venvs"
  )
  for d in "${dirs[@]}"; do
    run mkdir -p "$d"
  done
}

# -------------------------------
# APT (Desktop/User Packages)
install_apt_desktop() {
  log "Installing APT desktop/user packages..."
  run sudo apt update
  run sudo apt install -y \
    terminator \
    tmux \
    btop \
    neovim \
    autokey-gtk \
    folder-color \
    neofetch \
    copyq \
    tree \
    git \
    curl \
    wget \
    jq \
    ca-certificates \
    gnupg \
    software-properties-common
}

# -------------------------------
# Dev tools (kept separate)
install_dev_tools() {
  log "Installing dev tools..."
  run sudo apt update
  run sudo apt install -y \
    build-essential \
    cmake \
    ccache \
    clang \
    lld \
    gdb \
    ripgrep \
    fd-find \
    pkg-config \
    python3-venv \
    python3-pip
}

# -------------------------------
# Git config
setup_git() {
  log "Configuring Git..."

  if ! command -v git &>/dev/null; then
    run sudo apt update
    run sudo apt install -y git
  fi

  # NOTE: change email to your real one
  run git config --global user.name "Anirudh"
  run git config --global user.email "your@email.com"
  run git config --global init.defaultBranch main
  run git config --global pull.rebase false

  log "Git configured (name/email/defaultBranch)."
}

# -------------------------------
# Python venv: vision (no 'source', call venv pip directly)
setup_vision_venv() {
  log "Creating/updating Python venv: vision"

  run mkdir -p "$HOME/venvs"
  if [[ ! -d "$HOME/venvs/vision" ]]; then
    run python3 -m venv "$HOME/venvs/vision"
  else
    log "Venv already exists: $HOME/venvs/vision"
  fi

  run "$HOME/venvs/vision/bin/python" -m pip install --upgrade pip
  run "$HOME/venvs/vision/bin/pip" install opencv-python matplotlib seaborn pandas

  log "vision venv ready with: opencv-python, matplotlib, seaborn, pandas"
}

# -------------------------------
# PyTorch into vision venv (GPU-aware)
install_pytorch() {
  log "Installing PyTorch into vision venv..."

  if [[ ! -d "$HOME/venvs/vision" ]]; then
    die "vision venv not found. Run with --dev yes before --pytorch yes."
  fi

  if [[ "$GPU_TYPE" == "nvidia" ]]; then
    # CUDA wheels (cu118 is common; adjust if your CUDA stack differs)
    run "$HOME/venvs/vision/bin/pip" install torch torchvision torchaudio \
      --index-url https://download.pytorch.org/whl/cu118
    log "Installed PyTorch (CUDA wheels)."
  else
    run "$HOME/venvs/vision/bin/pip" install torch torchvision torchaudio
    log "Installed PyTorch (CPU wheels)."
  fi
}

# -------------------------------
# Docker (Install docker + fix perms + service on)
install_docker() {
  log "Installing Docker (engine + compose plugin)..."

  if command -v docker &>/dev/null; then
    log "Docker already installed: $(docker --version 2>/dev/null || echo "docker present")"
  else
    run sudo apt update
    run sudo apt install -y docker.io docker-compose-plugin
    log "Docker installed."
  fi

  run sudo systemctl enable docker || true
  run sudo systemctl start docker || true

  if getent group docker >/dev/null; then
    run sudo usermod -aG docker "$USER" || true
  else
    warn "docker group not found; creating it."
    run sudo groupadd docker || true
    run sudo usermod -aG docker "$USER" || true
  fi

  if docker info >/dev/null 2>&1; then
    log "Docker daemon reachable for current session."
  else
    warn "Docker installed, but current user may not have permission yet."
    warn "Log out and log back in (or run: newgrp docker) to use docker without sudo."
  fi

  run docker --version || true
  run docker compose version || true
}

# -------------------------------
# VS Code install + dev extensions + Docker helpers
install_vscode_extensions() {
  log "Installing VS Code + dev extensions (and Docker)..."

  # Ensure prerequisites
  run sudo apt update
  run sudo apt install -y curl gpg apt-transport-https ca-certificates

  # Install VS Code if missing
  if ! command -v code &>/dev/null; then
    log "VS Code not found. Installing VS Code..."

    run sudo install -d -m 0755 /etc/apt/keyrings

    # pipelines -> use bash -lc (no eval)
    run bash -lc \
      'curl -fsSL https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor | sudo tee /etc/apt/keyrings/packages.microsoft.gpg >/dev/null'
    run sudo chmod go+r /etc/apt/keyrings/packages.microsoft.gpg

    run bash -lc \
      'echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/packages.microsoft.gpg] https://packages.microsoft.com/repos/code stable main" | sudo tee /etc/apt/sources.list.d/vscode.list >/dev/null'

    run sudo apt update
    run sudo apt install -y code

    if ! command -v code &>/dev/null; then
      warn "VS Code install completed but 'code' still not on PATH. Open a new terminal and re-run."
      return 0
    fi

    log "VS Code installed: $(code --version | head -n 1 || true)"
  else
    log "VS Code already present: $(code --version | head -n 1 || true)"
  fi

  # Install Docker if missing
  if ! command -v docker &>/dev/null; then
    log "Docker not found. Installing Docker..."
    install_docker
  else
    log "Docker already present: $(docker --version 2>/dev/null || true)"
  fi

  # Your dev-useful extension set
  local exts=(
    # Productivity / UI
    "alefragnani.bookmarks"
    "vscode-icons-team.vscode-icons"
    "vscodevim.vim"
    "emilast.logfilehighlighter"

    # Python
    "ms-python.python"
    "ms-python.vscode-pylance"
    "ms-python.debugpy"
    "ms-python.vscode-python-envs"

    # C++ / CMake
    "ms-vscode.cpptools-themes"
    "ms-vscode.cmake-tools"
    "twxs.cmake"

    # Containers / Docker
    "ms-azuretools.vscode-docker"
    "docker.docker"
    "ms-azuretools.vscode-containers"
    "ms-vscode-remote.remote-containers"

    # Remote development (SSH/WSL/Remote pack)
    "ms-vscode-remote.remote-ssh"
    "ms-vscode-remote.remote-ssh-edit"
    "ms-vscode-remote.remote-wsl"
    "ms-vscode.remote-explorer"
    "ms-vscode.remote-server"
    "ms-vscode-remote.vscode-remote-extensionpack"

    # GitHub / Codespaces
    "github.codespaces"

    # Docs / misc dev helpers
    "shd101wyy.markdown-preview-enhanced"
    "tomoki1207.pdf"
  )

  for ext in "${exts[@]}"; do
    run code --install-extension "$ext" --force
  done

  log "VS Code extensions installed."

  if ! docker info >/dev/null 2>&1; then
    warn "Docker may require logout/login (or 'newgrp docker') before it works without sudo."
  fi
}

# -------------------------------
# Ollama (GPU-aware model pull)
install_ollama() {
  log "Installing Ollama..."

  if ! command -v curl &>/dev/null; then
    run sudo apt update
    run sudo apt install -y curl
  fi

  run bash -lc 'curl -fsSL https://ollama.com/install.sh | sh'

  if command -v ollama &>/dev/null; then
    if [[ "$GPU_TYPE" == "nvidia" ]]; then
      log "Pulling a general-purpose model (GPU box)."
      run ollama pull mistral
    else
      log "Pulling a lighter model (CPU box)."
      run ollama pull phi
    fi
  else
    warn "Ollama install did not put 'ollama' on PATH (unexpected)."
  fi
}

# -------------------------------
# Bashrc modifications (idempotent block)
modify_bashrc() {
  log "Modifying $HOME/.bashrc (NJ block)..."

  local bashrc="$HOME/.bashrc"
  local start="### NJ_SETUP_START ###"

  if grep -q "$start" "$bashrc" 2>/dev/null; then
    log "NJ bashrc block already present."
    return 0
  fi

  cat <<'EOF' >> "$bashrc"

### NJ_SETUP_START ###
# Quality of life
alias ll='ls -la'
alias ..='cd ..'
alias ...='cd ../..'

# Ensure local bin is available
export PATH="$HOME/.local/bin:$PATH"

# Git branch in prompt (only if available)
if [ -f /usr/lib/git-core/git-sh-prompt ]; then
  source /usr/lib/git-core/git-sh-prompt
fi

# Nice prompt: user: cwd (git-branch)
export PS1="${debian_chroot:+($debian_chroot)}\[\033[01;32m\]\u:\[\033[01;34m\]\w\[\033[00m\]\$(__git_ps1) \[\033[00m\]"
### NJ_SETUP_END ###
EOF

  log "$HOME/.bashrc updated. Open a new terminal or run: source $HOME/.bashrc"
}

# =========================================================
# Argument Parsing
# =========================================================
if [[ $# -eq 0 ]]; then
  usage
  exit 0
fi

while [[ $# -gt 0 ]]; do
  case "$1" in
    --install-all) need_value "$1" "${2:-}"; INSTALL_ALL="$2"; shift 2 ;;

    --ssh) need_value "$1" "${2:-}"; SSH_KEYS="$2"; shift 2 ;;
    --dirs) need_value "$1" "${2:-}"; DIRS="$2"; shift 2 ;;
    --apt) need_value "$1" "${2:-}"; APT_DESKTOP="$2"; shift 2 ;;
    --dev) need_value "$1" "${2:-}"; DEV="$2"; shift 2 ;;
    --bashrc) need_value "$1" "${2:-}"; BASHRC="$2"; shift 2 ;;
    --vscode) need_value "$1" "${2:-}"; VSCODE="$2"; shift 2 ;;
    --pytorch) need_value "$1" "${2:-}"; PYTORCH="$2"; shift 2 ;;
    --ollama) need_value "$1" "${2:-}"; OLLAMA="$2"; shift 2 ;;
    --dry-run) need_value "$1" "${2:-}"; DRY_RUN="$2"; shift 2 ;;

    -h|--help) usage; exit 0 ;;
    *) die "Unknown argument: $1 (use --help)" ;;
  esac
done

# Validate yes/no
validate_yesno "--install-all" "$INSTALL_ALL"
validate_yesno "--ssh" "$SSH_KEYS"
validate_yesno "--dirs" "$DIRS"
validate_yesno "--apt" "$APT_DESKTOP"
validate_yesno "--dev" "$DEV"
validate_yesno "--bashrc" "$BASHRC"
validate_yesno "--vscode" "$VSCODE"
validate_yesno "--pytorch" "$PYTORCH"
validate_yesno "--ollama" "$OLLAMA"
validate_yesno "--dry-run" "$DRY_RUN"

# If install-all: enable everything
if [[ "$INSTALL_ALL" == "yes" ]]; then
  # Note: do this AFTER validation to catch typos in any explicitly provided flags
  SSH_KEYS="yes"
  DIRS="yes"
  APT_DESKTOP="yes"
  DEV="yes"
  BASHRC="yes"
  VSCODE="yes"
  PYTORCH="yes"
  OLLAMA="yes"
fi

# =========================================================
# Execution
# =========================================================
print_banner
log "Logging to: $LOG_FILE"

detect_os
detect_gpu

if [[ "$INSTALL_ALL" == "yes" ]]; then
  log "--install-all enabled: running full setup."
fi

if [[ "$SSH_KEYS" == "yes" ]]; then setup_ssh_keys; fi
if [[ "$DIRS" == "yes" ]]; then create_dirs; fi
if [[ "$APT_DESKTOP" == "yes" ]]; then install_apt_desktop; fi

if [[ "$DEV" == "yes" ]]; then
  install_dev_tools
  setup_git
  setup_vision_venv
fi

if [[ "$PYTORCH" == "yes" ]]; then install_pytorch; fi
if [[ "$VSCODE" == "yes" ]]; then install_vscode_extensions; fi
if [[ "$OLLAMA" == "yes" ]]; then install_ollama; fi
if [[ "$BASHRC" == "yes" ]]; then modify_bashrc; fi

log "NJ Setup Complete."
warn "If you installed docker groups: log out and log back in."