bootstrap

#!/bin/bash
# -*- mode: bash -*-
set -eu -o pipefail

TARGET_OS=""

ROOT="$(git rev-parse --show-toplevel 2>/dev/null || (cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd))"
readonly ROOT
readonly TMPDIR=${TMPDIR:-/tmp}

readonly RISCV_DIR="$HOME/.local/riscv"
readonly STARSHIP_CONFIG="$HOME/.config/starship.toml"
readonly ZSH_CUSTOM_DIR="${ZSH_CUSTOM:-$HOME/.oh-my-zsh/custom}"

log() { printf "\033[1;32m==> %s\033[0m\n" "$*"; }
warn() { printf "\033[1;33m==> %s\033[0m\n" "$*"; }
err() { printf "\033[1;31m[!] %s\033[0m\n" "$*" >&2; }

INSTALL_SHELL_INTEGRATION=0

# Ensure command exists, install if missing
ensure() {
	local cmd="$1"
	local install_cmd="$2"
	if ! command -v "$cmd" >/dev/null 2>&1; then
		log "Installing $cmd..."
		eval "$install_cmd"
	else
		log "$cmd is already installed."
	fi
}

# Check if line exists in file, add if missing
affirm() {
	local line="$1"
	local file="$2"
	local description="${3:-Adding configuration to $file}"

	# Create the file if it doesn't exist
	touch "$file"

	# Check if the content already exists in the file
	if ! grep -qF "$line" "$file" 2>/dev/null; then
		log "$description"
		printf '%s\n' "$line" >>"$file"
	fi
}

# Run a command with sudo if not already root.
elevate() {
	if [ $# -ne 1 ]; then
		err "elevate expects a single command string argument"
		return 1
	fi
	local cmd="$1"
	if [ "${TARGET_OS:-}" = "macos" ] || [ "$(id -u)" -eq 0 ]; then
		bash -c "$cmd"
	else
		sudo bash -c "$cmd"
	fi
}

# Cross-platform sed.
sedi() {
	if [ "$TARGET_OS" = "macos" ]; then
		sed -i '' "$@"
	else
		sed -i "$@"
	fi
}

detect_os() {
	log "Detecting operating system..."
	case "$(uname -s)" in
	Linux) TARGET_OS=linux ;;
	Darwin) TARGET_OS=macos ;;
	*)
		err "Unsupported operating system: $(uname -s)"
		exit 1
		;;
	esac
	log "Detected $TARGET_OS"
}

get_rcfile_path() {
	local profile_path
	case "$SHELL" in
	*/zsh) profile_path="${ZDOTDIR:-$HOME}/.zshrc" ;;
	*/bash) profile_path="$HOME/.bashrc" ;;
	*)
		err "Could not detect current shell ($SHELL). Only zsh and bash are supported for rc configuration."
		return 1
		;;
	esac
	echo "$profile_path"
}

get_envrc_path() {
	echo "$ROOT/.envrc"
}

# Abstraction for package installation.
install_pkg() {
	log "Installing packages: $*"
	case "$TARGET_OS" in
	linux)
		if ! command -v apt-get >/dev/null 2>&1; then
			err "This script currently supports only Debian/Ubuntu-style systems (requires apt-get)."
			exit 1
		fi
		# Use $* (not $@) so this expands to a single string argument for elevate().
		elevate "apt-get update && apt-get install -y --no-install-recommends $*"
		;;
	macos)
		brew install "$@"
		;;
	*)
		err "Package installation not supported on $TARGET_OS"
		exit 1
		;;
	esac
}

# Install zsh plugin if directory doesn't exist
install_zsh_plugin() {
	local repo_url="$1"
	local plugin_dir="$2"
	local plugin_name
	plugin_name=$(basename "$plugin_dir")

	if [ ! -d "$plugin_dir" ]; then
		log "Installing $plugin_name..."
		git clone --depth=1 "$repo_url" "$plugin_dir" 2>/dev/null || {
			warn "Failed to clone $plugin_name"
			return 1
		}
	else
		log "$plugin_name is already installed."
	fi
}

setup_base_packages() {
	case "$TARGET_OS" in
	linux)
		install_pkg ca-certificates curl git sudo vim.tiny build-essential gh pkg-config libssl-dev fzf direnv
		elevate "update-alternatives --install /usr/bin/vim vim /usr/bin/vim.tiny 10"
		;;
	macos)
		install_pkg gh pkg-config openssl@3 coreutils fzf direnv
		;;
	esac
}

setup_direnv() {
	log "Setting up direnv..."

	if [ "${INSTALL_SHELL_INTEGRATION}" = "1" ]; then
		# Add direnv hook to shells
		affirm 'eval "$(direnv hook bash)"' "$HOME/.bashrc" "Adding direnv hook to bash"
		affirm 'eval "$(direnv hook zsh)"' "$HOME/.zshrc" "Adding direnv hook to zsh"
	fi

	# Create project .envrc file
	local envrc_file
	envrc_file=$(get_envrc_path)
	log "Creating project .envrc at $envrc_file"
	touch "$envrc_file"

	# Keep rustup from printing update/component check noise on every cargo invocation.
	affirm 'export RUSTUP_NO_UPDATE_CHECK=1' "$envrc_file" "Disabling rustup update checks for this repo"

	# Repo-local locale config (avoid mutating a user's global shell rc).
	affirm 'export LANG=en_US.UTF-8' "$envrc_file" "Setting LANG for this repo"
}

setup_zsh() {
	log "Setting up Zsh and Oh My Zsh..."
	if ! command -v zsh >/dev/null 2>&1; then
		install_pkg zsh
	fi

	if [[ ! -d "$HOME/.oh-my-zsh" ]]; then
		sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
	else
		log "Oh My Zsh is already installed."
	fi

	log "Installing Zsh plugins..."
	install_zsh_plugin "https://github.com/zsh-users/zsh-autosuggestions" "$ZSH_CUSTOM_DIR/plugins/zsh-autosuggestions"
	install_zsh_plugin "https://github.com/zsh-users/zsh-completions" "$ZSH_CUSTOM_DIR/plugins/zsh-completions"

	if [ "${INSTALL_SHELL_INTEGRATION}" = "1" ]; then
		touch "$HOME/.zshrc"
		log "Configuring Zsh plugins..."

		# Configure plugins
		sedi 's/^plugins.*)/plugins=(git fzf zsh-autosuggestions zsh-completions)/' "$HOME/.zshrc"

		# Fix locale issues (https://github.com/ohmyzsh/ohmyzsh/issues/7426)
		affirm 'export LANG=en_US.UTF-8' "$HOME/.zshrc"
	else
		local envrc_file
		envrc_file=$(get_envrc_path)
		affirm 'export LANG=en_US.UTF-8' "$envrc_file" "Setting LANG for this repo"
	fi
}

setup_bash() {
	if [ "${INSTALL_SHELL_INTEGRATION}" = "1" ]; then
		# Configure fzf for bash using modern integration
		affirm 'if command -v fzf >/dev/null 2>&1; then eval "$(fzf --bash)"; fi' "$HOME/.bashrc" "Adding fzf configuration"
	fi
}

setup_starship() {
	log "Installing and configuring Starship prompt..."
	ensure starship "curl -fsSL https://starship.rs/install.sh | sh -s -- --yes >/dev/null"

	mkdir -p "$(dirname "$STARSHIP_CONFIG")"
	if [ ! -f "$STARSHIP_CONFIG" ]; then
		cat >"$STARSHIP_CONFIG" <<'EOF'
format = """
$directory\
$git_branch$git_commit$git_state$git_metrics$git_status\
$line_break\
$character
"""
add_newline = false
EOF
	fi
	# Configure starship for current shell
	if [ "${INSTALL_SHELL_INTEGRATION}" = "1" ]; then
		affirm 'eval "$(starship init bash)"' "$HOME/.bashrc" "Configuring Starship for bash"
		affirm 'eval "$(starship init zsh)"' "$HOME/.zshrc" "Configuring Starship for zsh"
	fi
}

setup_submodules() {
	log "Initializing and updating git submodules..."
	git submodule update --init --recursive
}

setup_git_hooks() {
	log "Configuring Git hooks..."
	# git config --local core.hooksPath ".ghk"
}

setup_rust() {
	if ! command -v rustup >/dev/null 2>&1; then
		log "Setting up Rust toolchain..."

		ensure rustup "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --no-update-default-toolchain"
		# Add cargo to project .envrc
		local envrc_file
		envrc_file=$(get_envrc_path)
		affirm 'export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:$PATH"' "$envrc_file" "Adding Cargo to PATH"
	else
		log "rustup is already installed."
	fi
	# Source the environment to make cargo available for the rest of the script
	. "${CARGO_HOME:-$HOME/.cargo}/env"
}

setup_cargo() {
	log "Installing essential cargo tools (just, cargo-nextest, taplo)..."
	if ! command -v cargo >/dev/null 2>&1; then
		rustup toolchain install stable -q
	fi
	ensure just "cargo install --locked just"
	ensure cargo-nextest "cargo install --locked cargo-nextest"
	ensure taplo "cargo install --locked taplo-cli"
}

setup_shfmt() {
	log "Installing shfmt (shell formatter)..."
	install_pkg shfmt
}

setup_deno() {
	ensure deno "curl -fsSL https://deno.land/install.sh | sh -s -- --yes --no-modify-path"
	# Add deno to project .envrc
	local envrc_file
	envrc_file=$(get_envrc_path)
	affirm 'export PATH="$HOME/.deno/bin:$PATH"' "$envrc_file" "Adding Deno to PATH"
}

setup_act() {
	log "Installing act (GitHub Actions runner)..."
	case "$TARGET_OS" in
	linux)
		if ! command -v act >/dev/null 2>&1; then
			log "Downloading and installing act..."
			curl --proto '=https' --tlsv1.2 -sSf https://raw.githubusercontent.com/nektos/act/master/install.sh | bash -s -- -b "$HOME/.local/bin"
		else
			log "act is already installed."
		fi
		;;
	macos)
		install_pkg act
		;;
	esac
}

setup_gh_auth() {
	log "Checking GitHub authentication..."

	# In CI/non-interactive contexts we never want to prompt.
	# If a token is available, assume authentication is handled by the environment.
	if [ "${CI:-}" = "true" ] || [ "${GITHUB_ACTIONS:-}" = "true" ] || [ ! -t 0 ]; then
		if [ -n "${GH_TOKEN:-}" ] || [ -n "${GITHUB_TOKEN:-}" ]; then
			log "CI/non-interactive environment detected; token present, skipping 'gh auth login'."
			return 0
		fi
		warn "CI/non-interactive environment detected and no GH token env var found; skipping 'gh auth login'."
		return 0
	fi

	if ! gh auth status >/dev/null 2>&1; then
		warn "You are not logged into GitHub. Please run 'gh auth login'."
		gh auth login
	else
		log "GitHub authentication confirmed."
	fi
}

setup_riscv_toolchain() {
	log "Installing RISC-V toolchain..."
	case "$TARGET_OS" in
	linux)
		install_pkg device-tree-compiler gcc-riscv64-unknown-elf binutils-riscv64-unknown-elf picolibc-riscv64-unknown-elf
		;;
	macos)
		warn "RISC-V toolchain on macOS requires manual installation via Homebrew."
		warn "Please run: brew tap riscv/riscv && brew install riscv-gnu-toolchain"
		;;
	esac
}

setup_spike() {
	log "Installing Spike RISC-V simulator..."
	local spike_archive="/tmp/spike.tar.gz"
	trap 'rm -f /tmp/spike.tar.gz' EXIT

	case "$TARGET_OS" in
	linux)
		if ! gh release download --clobber --repo LayerZero-Labs/ZeroOS spike-1.1.1 --pattern "spike-1.1.1-$(uname -s)-$(uname -m).tar.gz" -O "$spike_archive"; then
			err "Failed to download Spike release. Check if the release exists and you have network access."
			return 1
		fi
		elevate "mkdir -p $RISCV_DIR && tar -xzf $spike_archive -C $RISCV_DIR"
		;;
	macos)
		install_pkg riscv-isa-sim
		;;
	esac

	# Add project-specific environment variables to current shell's RC file
	local rc_file
	rc_file=$(get_envrc_path) || return
	affirm "export RISCV=\"${RISCV_DIR/$HOME/\$HOME}\"" "$rc_file" "Setting RISCV environment variable"
	affirm "export PATH=\"\$RISCV/bin:\$PATH\"" "$rc_file" "Adding RISC-V toolchain to PATH"
}

setup_shell_rc() {
	log "Configuring shell environment..."

	# Add general PATH exports to current shell's RC file
	local rc_file
	rc_file=$(get_envrc_path) || return
	affirm 'export PATH="$HOME/.local/bin:$PATH"' "$rc_file" "Adding local bin directory to PATH"

	log "Shell environment configured."
}

main() {
	while [ $# -gt 0 ]; do
		case "$1" in
		--install-shell-integration)
			INSTALL_SHELL_INTEGRATION=1
			shift
			;;
		-h | --help)
			cat <<'EOF'
Usage: ./bootstrap [--install-shell-integration]

By default, this script does NOT modify ~/.bashrc or ~/.zshrc.
Use --install-shell-integration to update shell rc files with direnv/starship hooks and zsh plugins.
EOF
			exit 0
			;;
		*)
			err "Unknown argument: $1"
			exit 1
			;;
		esac
	done

	detect_os

	cd "$ROOT"
	git config --global --add safe.directory "$ROOT"

	setup_base_packages
	setup_direnv
	setup_starship
	setup_zsh
	setup_bash
	setup_act
	setup_submodules
	setup_git_hooks
	setup_rust
	setup_cargo
	setup_shfmt
	setup_deno
	setup_gh_auth
	setup_riscv_toolchain
	setup_spike
	setup_shell_rc

	# Allow direnv to load the .envrc file
	log "Allowing direnv to load .envrc..."
	direnv allow

	log "Bootstrap complete!"
	if [ "${INSTALL_SHELL_INTEGRATION}" != "1" ]; then
		warn "Note: this script did NOT modify ~/.bashrc or ~/.zshrc."
		warn "If you want shell integration, re-run with '--install-shell-integration',"
		warn "or add the following to your rc file(s):"
		cat <<'EOF'

## bash (~/.bashrc)
eval "$(direnv hook bash)"
eval "$(starship init bash)"
if command -v fzf >/dev/null 2>&1; then eval "$(fzf --bash)"; fi

## zsh (~/.zshrc)
eval "$(direnv hook zsh)"
eval "$(starship init zsh)"
plugins=(git fzf zsh-autosuggestions zsh-completions)

EOF
	else
		log "Shell integration installed into ~/.bashrc and/or ~/.zshrc"
	fi
	log "Please restart your shell or run 'direnv allow' to activate the environment."
	log "  or run 'exec \$SHELL -l' to reload your shell."
}

main "$@"