install.m4

#!/bin/bash
# shellcheck disable=SC2317
# shellcheck disable=SC2154

# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

# m4_ignore(
echo "This is just a script template, not the script (yet) - pass it to 'argbash' to fix this." >&2
exit 11 #)
# ARG_HELP([A one-stop-shop for installing the Tenstorrent stack])
# ARG_VERSION([echo "__INSTALLER_DEVELOPMENT_BUILD__"])
# ========================= Boolean Arguments =========================
# ARG_OPTIONAL_BOOLEAN([install-kmd],,[Kernel-Mode-Driver installation],[on])
# ARG_OPTIONAL_BOOLEAN([install-hugepages],,[Configure HugePages],[on])
# ARG_OPTIONAL_BOOLEAN([install-podman],,[Install Podman],[on])
# ARG_OPTIONAL_BOOLEAN([install-podman-docker],,[Install podman-docker shim (disable to keep docker)],[off])
# ARG_OPTIONAL_BOOLEAN([install-metalium-container],,[Download and install Metalium container],[on])
# ARG_OPTIONAL_BOOLEAN([install-tt-flash],,[Install tt-flash for updating device firmware],[on])
# ARG_OPTIONAL_BOOLEAN([install-tt-smi],,[Install tt-smi for device monitoring],[on])
# ARG_OPTIONAL_BOOLEAN([install-tt-topology],,[Install tt-topology (Wormhole only)],[off])
# ARG_OPTIONAL_BOOLEAN([install-sfpi],,[Install SFPI],[on])
# ARG_OPTIONAL_BOOLEAN([install-inference-server],,[Install tt-inference-server],[on])
# ARG_OPTIONAL_BOOLEAN([install-studio],,[Install tt-studio],[on])

# =========================  Podman Metalium Arguments =========================
# ARG_OPTIONAL_SINGLE([metalium-image-url],,[Container image URL to pull/run],[ghcr.io/tenstorrent/tt-metal/tt-metalium-ubuntu-22.04-release-amd64])
# ARG_OPTIONAL_SINGLE([metalium-image-tag],,[Tag (version) of the Metalium image],[latest-rc])
# ARG_OPTIONAL_SINGLE([podman-metalium-script-dir],,[Directory where the helper wrapper will be written],["$HOME/.local/bin"])
# ARG_OPTIONAL_SINGLE([podman-metalium-script-name],,[Name of the helper wrapper script],["tt-metalium"])
# ARG_OPTIONAL_BOOLEAN([install-metalium-models-container],,[Install additional TT-Metalium container for running model demos],[off])

# ========================= String Arguments =========================
# ARG_OPTIONAL_SINGLE([python-choice],,[Python setup strategy: active-venv, new-venv, system-python, pipx],[new-venv])
# ARG_OPTIONAL_BOOLEAN([use-uv],,[Use uv instead of pip for Python package installation],[off])
# ARG_OPTIONAL_SINGLE([reboot-option],,[Reboot policy after install: ask, never, always],[ask])
# ARG_OPTIONAL_SINGLE([update-firmware],,[Update TT device firmware: on, off, force],[on])
# ARG_OPTIONAL_SINGLE([github-token],,[Optional GitHub API auth token],[])

# ========================= Version Arguments =========================
# ARG_OPTIONAL_SINGLE([kmd-version],,[Specific version of TT-KMD to install],[])
# ARG_OPTIONAL_SINGLE([fw-version],,[Specific version of firmware to install],[])
# ARG_OPTIONAL_SINGLE([systools-version],,[Specific version of system tools to install],[])
# ARG_OPTIONAL_SINGLE([smi-version],,[Specific version of tt-smi to install],[])
# ARG_OPTIONAL_SINGLE([flash-version],,[Specific version of tt-flash to install],[])
# ARG_OPTIONAL_SINGLE([topology-version],,[Specific version of tt-topology to install],[])
# ARG_OPTIONAL_SINGLE([sfpi-version],,[Specific version of SFPI to install],[])

# ========================= Path Arguments =========================
# ARG_OPTIONAL_SINGLE([new-venv-location],,[Path for new Python virtual environment],[$HOME/.tenstorrent-venv])

# ========================= Mode Arguments =========================
# ARG_OPTIONAL_BOOLEAN([mode-container],,[Enable container mode (skips KMD, HugePages, and SFPI, never reboots)],[off])
# ARG_OPTIONAL_BOOLEAN([mode-non-interactive],,[Enable non-interactive mode (no user prompts)],[off])
# ARG_OPTIONAL_BOOLEAN([verbose],,[Enable verbose output for debugging])

# ARGBASH_GO

# [ <-- needed because of Argbash

# Logo
# Credit: figlet font slant by Glenn Chappell
LOGO=$(cat << "EOF"
   __                  __                             __
  / /____  ____  _____/ /_____  _____________  ____  / /_
 / __/ _ \/ __ \/ ___/ __/ __ \/ ___/ ___/ _ \/ __ \/ __/
/ /_/  __/ / / (__  ) /_/ /_/ / /  / /  /  __/ / / / /_
\__/\___/_/ /_/____/\__/\____/_/  /_/   \___/_/ /_/\__/
EOF
)

# If container mode is enabled, disable KMD, HugePages, and SFPI
# shellcheck disable=SC2154
if [[ "${_arg_mode_container}" = "on" ]]; then
	_arg_install_kmd="off"
	_arg_install_hugepages="off" # Both KMD and HugePages must live on the host kernel
	_arg_install_podman="off" # No podman in podman
	_arg_install_sfpi="off"
	_arg_reboot_option="never" # Do not reboot
fi

PIPX_ENSUREPATH_EXTRAS="${TT_PIPX_ENSUREPATH_EXTRAS:- }"
PIPX_INSTALL_EXTRAS="${TT_PIPX_INSTALL_EXTRAS:- }"

# ========================= Main Script =========================

# Create working directory
TMP_DIR_TEMPLATE="tenstorrent_install_XXXXXX"
# Use mktemp to get a temporary directory
WORKDIR=$(mktemp -d -p /tmp "${TMP_DIR_TEMPLATE}")

# Initialize logging
LOG_FILE="${WORKDIR}/install.log"
# Redirect stdout to the logfile.
# Removes color codes and prepends the date
exec > >( \
		tee >( \
				stdbuf -o0 \
						sed 's/\x1B\[[0-9;]*[A-Za-z]//g' | \
						xargs -d '\n' -I {} date '+[%F %T] {}' \
				> "${LOG_FILE}" \
				) \
		)
exec 2>&1

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# argbash workaround: close square brackets ]]]]]

# log messages to terminal (with color)
log() {
	local msg="[INFO] $1"
	echo -e "${GREEN}${msg}${NC}"  # Color output to terminal
}

# log errors
error() {
	local msg="[ERROR] $1"
	echo -e "${RED}${msg}${NC}"
}

# log an error and then exit
error_exit() {
    error "$1"
    exit 1
}

# log warnings
warn() {
	local msg="[WARNING] $1"
	echo -e "${YELLOW}${msg}${NC}"
}

check_has_sudo_perms() {
	if ! sudo true; then
		error "Cannot use sudo, exiting..."
		exit 1
	fi
}

detect_distro() {
	# shellcheck disable=SC1091 # Always present
	if [[ -f /etc/os-release ]]; then
		. /etc/os-release
		DISTRO_ID=${ID}
		case ${DISTRO_ID} in
			ubuntu|debian|fedora|rhel|centos)
				;; # It's a known distro, do nothing.
			*)
				# Could be a derivative distribution, check ID_LIKE
				if [[ -n "${ID_LIKE:-}" ]]; then
					# ID_LIKE can be a space-separated list. Check each of them.
					for id_like_distro in ${ID_LIKE}; do
						case ${id_like_distro} in
							ubuntu|debian|fedora|rhel)
								DISTRO_ID=${id_like_distro}
								break # Use the first one found.
								;;
						esac
					done
				fi
				;;
		esac
		# Set package manager based on distribution
		case "${DISTRO_ID}" in
			"ubuntu"|"debian")
				PKG_MANAGER="apt-get"
				;;
			"fedora"|"rhel"|"centos")
				PKG_MANAGER="dnf"
				;;
			*)
				error "Unsupported distribution: ${DISTRO_ID}"
				exit 1
				;;
		esac
	else
		error "Cannot detect Linux distribution"
		exit 1
	fi
}

# Function to verify download
verify_download() {
	local file=$1
	if [[ ! -f "${file}" ]]; then
		error "Download failed: ${file} not found"
		exit 1
	fi
}

# Function to prompt for yes/no
confirm() {
	# In non-interactive mode, always return true
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		return 0
	fi

	while true; do
		read -rp "$1 [Y/n] " yn
		case ${yn} in
			[Nn]* ) echo && return 1;;
			[Yy]* | "" ) echo && return 0;;
			* ) echo "Please answer yes or no.";;
		esac
	done
}

set_non_interactive_defaults() {
	if [[ "${_arg_mode_non_interactive}" = "on" ]] && [[ "${_arg_reboot_option}" = "ask" ]]; then
		_arg_reboot_option="never" # Do not reboot in non-interactive mode
	fi
}

maybe_enable_default_mode() {
	# Respect already-set non-interactive mode
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		return
	fi

	log "Would you like to proceed with the default installation?"
	log "Selecting yes enables non-interactive mode and continues with default options."
	if confirm "Proceed with default installation (recommended for most users)"; then
		_arg_mode_non_interactive="on"
		set_non_interactive_defaults
		log "Default installation selected. Continuing in non-interactive mode."
	fi
}

# Function to check if uv is installed
check_uv_installed() {
	command -v uv &> /dev/null
}

# Get Python installation choice interactively or use default
get_python_choice() {
	PYTHON_CHOICE="${_arg_python_choice}"

	# In non-interactive mode, use the provided argument
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		log "Non-interactive mode, using Python installation method: ${PYTHON_CHOICE}"
	else
		log "How would you like to install Python packages?"
		# Interactive mode - show current choice and allow override
		while true; do
			echo "1) active-venv: Use the active virtual environment"
			echo "2) new-venv: [DEFAULT] Create a new Python virtual environment (venv) at ${_arg_new_venv_location}"
			echo "3) system-python: Use the system pathing, available for multiple users. *** NOT RECOMMENDED UNLESS YOU ARE SURE ***"
			echo "4) pipx: Use pipx for isolated package installation"
			read -rp "Enter your choice (1-4) or press enter for default (${_arg_python_choice}): " user_choice
			echo # newline

			# If user provided no value, use default and exit
			if [[ -z "${user_choice}" ]]; then
				break
			fi

			# Process user choice
			case "${user_choice}" in
				1|active-venv)
					PYTHON_CHOICE="active-venv"
					break
					;;
				2|new-venv)
					PYTHON_CHOICE="new-venv"
					break
					;;
				3|system-python)
					PYTHON_CHOICE="system-python"
					break
					;;
				4|pipx)
					PYTHON_CHOICE="pipx"
					break
					;;
				*)
					warn "Invalid choice '${user_choice}'. Please try again."
					;;
			esac
		done
	fi

	# Validate --use-uv flag
	if [[ "${_arg_use_uv}" = "on" ]]; then
		if ! check_uv_installed; then
			error "uv is not installed!"
			error_exit "Please install uv first: curl -LsSf https://astral.sh/uv/install.sh | sh"
		fi
		if [[ "${PYTHON_CHOICE}" = "pipx" ]]; then
			warn "--use-uv is not compatible with pipx, ignoring --use-uv flag"
			_arg_use_uv="off"
		else
			log "Using uv instead of pip for package installation"
		fi
	fi

	# Set up Python environment based on choice
	case ${PYTHON_CHOICE} in
		"active-venv")
			if [[ -z "${VIRTUAL_ENV:-}" ]]; then
				error "No active virtual environment detected!"
				error_exit "Please activate your virtual environment first and try again"
			fi
			log "Using active virtual environment: ${VIRTUAL_ENV}"
			INSTALLED_IN_VENV=0
			if [[ "${_arg_use_uv}" = "on" ]]; then
				PYTHON_INSTALL_CMD="uv pip install"
			else
				PYTHON_INSTALL_CMD="pip install"
			fi
			;;
		"system-python")
			log "Using system pathing"
			INSTALLED_IN_VENV=1
			# Check Python version to determine if --break-system-packages is needed (Python 3.11+)
			PYTHON_VERSION_MINOR=$(python3 -c "import sys; print(f'{sys.version_info.minor}')")
			if [[ "${_arg_use_uv}" = "on" ]]; then
				if [[ ${PYTHON_VERSION_MINOR} -gt 10 ]]; then
					PYTHON_INSTALL_CMD="uv pip install --system --break-system-packages"
				else
					PYTHON_INSTALL_CMD="uv pip install --system"
				fi
			else
				if [[ ${PYTHON_VERSION_MINOR} -gt 10 ]]; then
					PYTHON_INSTALL_CMD="pip install --break-system-packages"
				else
					PYTHON_INSTALL_CMD="pip install"
				fi
			fi
			;;
		"pipx")
			log "Using pipx for isolated package installation"
			# adding quotes around PIPX_ENSUREPATH_EXTRAS means they won't be
			# interpreted, which is exactly what we want them to be
			# shellcheck disable=2086
			pipx ensurepath ${PIPX_ENSUREPATH_EXTRAS}
			# Enable the pipx path in this shell session
			export PATH="${PATH}:${HOME}/.local/bin/"
			INSTALLED_IN_VENV=1
			PYTHON_INSTALL_CMD="pipx install ${PIPX_INSTALL_EXTRAS}"
			;;
		"new-venv"|*)
			log "Setting up new Python virtual environment"
			python3 -m venv "${_arg_new_venv_location}"
			# shellcheck disable=SC1091 # Must exist after previous command
			source "${_arg_new_venv_location}/bin/activate"
			INSTALLED_IN_VENV=0
			if [[ "${_arg_use_uv}" = "on" ]]; then
				PYTHON_INSTALL_CMD="uv pip install"
			else
				PYTHON_INSTALL_CMD="pip install"
			fi
			;;
	esac

}

# Function to check if Podman is installed
check_podman_installed() {
	command -v podman &> /dev/null
}

# Function to setup rootless Podman
setup_rootless_podman() {
	log "Configuring rootless Podman"
	# Add GUIDs/UIDs for rootless Podman
	# See https://github.com/containers/podman/blob/main/docs/tutorials/rootless_tutorial.md
	sudo usermod --add-subgids 10000-75535 "$(whoami)"
	sudo usermod --add-subuids 10000-75535 "$(whoami)"
}

# Install Podman Metalium container
install_podman_metalium() {
	log "Installing Metalium via Podman"

	# Create wrapper script directory
	mkdir -p "${_arg_podman_metalium_script_dir}" || error_exit "Failed to create script directory"

	# Create wrapper script
	log "Creating wrapper script..."
	cat > "${_arg_podman_metalium_script_dir}/${_arg_podman_metalium_script_name}" << EOF
#!/bin/bash
# Wrapper script for tt-metalium using Podman

# Image configuration
METALIUM_IMAGE="${_arg_metalium_image_url}:${_arg_metalium_image_tag}"

# Run the command using Podman

podman run --rm -it \\
  --privileged \\
  --log-driver none \\
  --volume=/dev/hugepages-1G:/dev/hugepages-1G \\
  --volume=\${HOME}:/home/user \\
  --device=/dev/tenstorrent:/dev/tenstorrent \\
  --workdir=/home/user \\
  --env=DISPLAY=\${DISPLAY} \\
  --env=HOME=/home/user \\
  --env=TERM=\${TERM:-xterm-256color} \\
  --network=host \\
  --security-opt label=disable \\
  --entrypoint /bin/bash \\
  \${METALIUM_IMAGE} "\$@"
EOF

	# Make the script executable
	chmod +x "${_arg_podman_metalium_script_dir}/${_arg_podman_metalium_script_name}" || error_exit "Failed to make script executable"

	# Check if the directory is in PATH
	if [[ ":${PATH}:" != *":${_arg_podman_metalium_script_dir}:"* ]]; then
		warn "${_arg_podman_metalium_script_dir} is not in your PATH."
		warn "A restart may fix this, or you may need to update your shell RC"
	fi

	# Pull the image
	log "Pulling the tt-metalium image (this may take a while)..."
	podman pull "${_arg_metalium_image_url}:${_arg_metalium_image_tag}" || error "Failed to pull image"

	log "Metalium installation completed"
	return 0
}

# Install Podman Metalium "models" container
install_podman_metalium_models() {
	log "Installing Metalium Models Container via Podman"
	local PODMAN_METALIUM_MODELS_SCRIPT_DIR="${HOME}/.local/bin"
	local PODMAN_METALIUM_MODELS_SCRIPT_NAME="tt-metalium-models"
	local METALIUM_MODELS_IMAGE_TAG="latest-rc"
	local METALIUM_MODELS_IMAGE_URL="ghcr.io/tenstorrent/tt-metal/tt-metalium-ubuntu-22.04-release-models-amd64"

	# Create wrapper script directory
	mkdir -p "${PODMAN_METALIUM_MODELS_SCRIPT_DIR}" || error_exit "Failed to create script directory"

	# Create wrapper script
	log "Creating wrapper script..."
	cat > "${PODMAN_METALIUM_MODELS_SCRIPT_DIR}/${PODMAN_METALIUM_MODELS_SCRIPT_NAME}" << EOF
#!/bin/bash
# Wrapper script for tt-metalium-models using Podman

echo "================================================================================"
echo "NOTE: This container tool for tt-metalium is meant to enable users to try out"
echo "      demos, and is not meant for production use. This container is liable to"
echo "      to change at anytime."
echo ""
echo "      For more information see https://github.com/tenstorrent/tt-metal/issues/25602"
echo "================================================================================"

# Image configuration
METALIUM_IMAGE="${METALIUM_MODELS_IMAGE_URL}:${METALIUM_MODELS_IMAGE_TAG}"

# Run the command using Podman
#
# Explaining some changes:
#  removal of --volume=\${HOME}:/home/user \\: the user in the upstream monster
#  container is user, and we put the source code in that user's directory, so
#  this would override it
#
#  removal of --workdir=/home/user \\: not super needed, but it's nice for
#  people to just be in the source code, ready to go
#
#  addition of --entrypoint /bin/bash: The current upstream container needs to
#  override the entrypoint. Why not just corral users into /bin/bash?
podman run --rm -it \\
  --privileged \\
  --log-driver none \\
  --volume=/dev/hugepages-1G:/dev/hugepages-1G \\
  --device=/dev/tenstorrent:/dev/tenstorrent \\
  --env=DISPLAY=\${DISPLAY} \\
  --env=HOME=/home/user \\
  --env=TERM=\${TERM:-xterm-256color} \\
  --network=host \\
  --security-opt label=disable \\
  --entrypoint /bin/bash \\
  \${METALIUM_IMAGE} "\$@"
EOF

	# Make the script executable
	chmod +x "${PODMAN_METALIUM_MODELS_SCRIPT_DIR}/${PODMAN_METALIUM_MODELS_SCRIPT_NAME}" || error_exit "Failed to make script executable"

	# Check if the directory is in PATH
	if [[ ":${PATH}:" != *":${PODMAN_METALIUM_MODELS_SCRIPT_DIR}:"* ]]; then
		warn "${PODMAN_METALIUM_MODELS_SCRIPT_DIR} is not in your PATH."
		warn "A restart may fix this, or you may need to update your shell RC"
	fi

	# Pull the image
	log "Pulling the tt-metalium-models image (this may take a while)..."
	podman pull "${METALIUM_MODELS_IMAGE_URL}:${METALIUM_MODELS_IMAGE_TAG}" || error "Failed to pull image"

	log "Metalium Models installation completed"
	return 0
}

get_podman_metalium_choice() {
	# In non-interactive mode, use the provided arguments
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		log "Non-interactive mode, using Podman Metalium installation preference: ${_arg_install_metalium_container}"
		log "Non-interactive mode, using Metalium Models installation preference: ${_arg_install_metalium_models_container}"
		return
	fi
	# Only ask if Podman is installed or will be installed
	if [[ "${_arg_install_podman}" = "on" ]] || check_podman_installed; then
		# Interactive mode - allow override
		log "Would you like to install the TT-Metalium slim container?"
		log "This container is appropriate if you only need to use TT-NN"
		if confirm "Install Metalium"; then
			_arg_install_metalium_container="on"
		else
			_arg_install_metalium_container="off"
		fi
	else
		# Podman won't be installed, so don't install Metalium
		_arg_install_metalium_container="off"
		warn "Podman is not and will not be installed, skipping Podman Metalium installation"
	fi
	# Only ask if Podman is installed or will be installed
	if [[ "${_arg_install_podman}" = "on" ]] || check_podman_installed; then
		# Interactive mode - allow override
		log "Would you like to install the TT-Metalium Model Demos container?"
		log "This container is best for users who need more TT-Metalium functionality, such as running prebuilt models, but it's large (8GB)"
		if confirm "Install Metalium Models"; then
			_arg_install_metalium_models_container="on"
		else
			_arg_install_metalium_models_container="off"
		fi
	else
		# Podman won't be installed, so don't install Metalium
		_arg_install_metalium_models_container="off"
		warn "Podman is not and will not be installed, skipping Podman Metalium Models installation"
	fi

	# Disable Podman if both Metalium containers are disabled
	if [[ "${_arg_install_metalium_container}" = "off" ]] && [[ "${_arg_install_metalium_models_container}" = "off" ]]; then
		_arg_install_podman="off"
	fi
}

get_inference_server_choice() {
	# In non-interactive mode, use the provided argument
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		log "Non-interactive mode, using tt-inference-server installation preference: ${_arg_install_inference_server}"
		return
	fi

	# Interactive mode - allow override
	log "Would you like to install tt-inference-server?"
	log "This will clone the inference server repository to ~/.local/lib and create a wrapper script"
	if confirm "Install tt-inference-server"; then
		_arg_install_inference_server="on"
	else
		_arg_install_inference_server="off"
	fi
}

get_studio_choice() {
	# In non-interactive mode, use the provided argument
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		log "Non-interactive mode, using tt-studio installation preference: ${_arg_install_studio}"
		return
	fi

	# Interactive mode - allow override
	log "Would you like to install tt-studio?"
	log "This will clone the tt-studio repository to ~/.local/lib and create a wrapper script"
	if confirm "Install tt-studio"; then
		_arg_install_studio="on"
	else
		_arg_install_studio="off"
	fi
}

# Generic function to fetch latest version from any GitHub repository
# Usage: fetch_latest_version <repo> <prefix_to_remove>
# Returns: version string with prefix removed, or exits with error code
fetch_latest_version() {
	local repo="$1"
	local prefix_to_remove="${2:-}"

	if ! command -v jq &> /dev/null; then
		echo "Error: JQ is not installed!" >&2
		return 1  # jq not installed
	fi

	local response
	local response_headers
	local response_body
	local latest_version

	# Curl options
	# We always suppress connect headers (fixes issues with systems using proxies)
	# -D - dumps the headers to stdout
	curl_opts=(--suppress-connect-headers -D -)

	# SC is worried this might not exist, but argbash guarantees it will
    # shellcheck disable=SC2154
	if [[ "${_arg_verbose}" = "on" ]]; then
		curl_opts+=(-v)
	else
		curl_opts+=(-s -S)
	fi

	if [[ -n "${_arg_github_token}" ]]; then
		curl_opts+=(-H "Authorization: token ${_arg_github_token}")
	fi

	response=$(curl "${curl_opts[@]}" \
		https://api.github.com/repos/"${repo}"/releases/latest)

	# Split at the first blank line
	response_headers=$(echo "${response}" | sed '/^\r*$/,$d')
	response_body=$(echo "${response}" | sed '1,/^\r*$/d')

	if [[ "${_arg_verbose}" = "on" ]]; then
		echo "=== GitHub API Response Headers ===" >&2
		echo "${response_headers}" >&2
		echo "=== GitHub API Response Body ===" >&2
		echo "${response_body}" >&2
		echo "===================================" >&2
	fi

	# Check for GitHub API rate limit
	if echo "${response_headers}" | grep -qi "x-ratelimit-remaining: 0"; then
		echo "Error: GitHub API Rate Limit exceeded" >&2
		return 2  # GitHub API rate limit exceeded
	fi

	# Check if response body is valid JSON
	if ! echo "${response_body}" | jq . >/dev/null 2>&1; then
		echo "Error: Got invalid JSON from GitHub API" >&2
		return 3  # Invalid JSON response
	fi

	latest_version=$(echo "${response_body}" | jq -r '.tag_name' 2>/dev/null)

	# Check if we got a valid tag_name
	if [[ -z "${latest_version}" || "${latest_version}" == "null" ]]; then
		echo "Error: No tag name found in API response" >&2
		return 4  # No tag_name found
	fi

	# Remove prefix if specified
	if [[ -n "${prefix_to_remove}" ]]; then
		echo "${latest_version#"${prefix_to_remove}"}"
	else
		echo "${latest_version}"
	fi

	return 0
}

install_tt_repos () {
	log "Installing TT repositories to your distribution package manager"
	case "${DISTRO_ID}" in
		"ubuntu")
			# Add the apt listing
			# shellcheck disable=2002
			echo "deb [signed-by=/etc/apt/keyrings/tt-pkg-key.asc] https://ppa.tenstorrent.com/ubuntu/ $( cat /etc/os-release | grep "^VERSION_CODENAME=" | sed 's/^VERSION_CODENAME=//' ) main" | sudo tee /etc/apt/sources.list.d/tenstorrent.list > /dev/null

			# Setup the keyring
			sudo mkdir -p /etc/apt/keyrings; sudo chmod 755 /etc/apt/keyrings

			# Download the key
			sudo wget -O /etc/apt/keyrings/tt-pkg-key.asc https://ppa.tenstorrent.com/tt-pkg-key.asc

			sudo apt-get update
			;;
		"debian")
			# Add the apt listing
			# shellcheck disable=2002
			echo "deb [signed-by=/etc/apt/keyrings/tt-pkg-key.asc] https://ppa.tenstorrent.com/debian/ $( cat /etc/os-release | grep "^VERSION_CODENAME=" | sed 's/^VERSION_CODENAME=//' ) main" | sudo tee /etc/apt/sources.list.d/tenstorrent.list > /dev/null

			# Setup the keyring
			sudo mkdir -p /etc/apt/keyrings; sudo chmod 755 /etc/apt/keyrings

			# Download the key
			sudo wget -O /etc/apt/keyrings/tt-pkg-key.asc https://ppa.tenstorrent.com/tt-pkg-key.asc

			sudo apt-get update
			;;
		"fedora")
			sudo bash -c 'cat > /etc/yum.repos.d/tenstorrent.repo << EOF
[Tenstorrent]
name=Tenstorrent
baseurl=https://ppa.tenstorrent.com/fedora
enabled=1
gpgcheck=1
gpgkey=http://ppa.tenstorrent.com/tt-pkg-key.asc
EOF'
			;;
		"rhel"|"centos")
			warn "RHEL and CentOS are not officially supported. Using Fedora repos."
			sudo bash -c 'cat > /etc/yum.repos.d/tenstorrent.repo << EOF
[Tenstorrent]
name=Tenstorrent
baseurl=https://ppa.tenstorrent.com/fedora
enabled=1
gpgcheck=1
gpgkey=http://ppa.tenstorrent.com/tt-pkg-key.asc
EOF'
			;;
		*)
			error_exit "Unsupported distro: ${DISTRO_ID}"
			;;
	esac
}

install_inference_server () {
	log "Installing tt-inference-server"
	local INFERENCE_SERVER_LIB_DIR="${HOME}/.local/lib"
	local INFERENCE_SERVER_BIN_DIR="${HOME}/.local/bin"
	local INFERENCE_SERVER_SCRIPT_NAME="tt-inference-server"
	local INFERENCE_SERVER_REPO_URL="https://github.com/tenstorrent/tt-inference-server.git"

	# Create directories
	mkdir -p "${INFERENCE_SERVER_LIB_DIR}" || error_exit "Failed to create library directory"
	mkdir -p "${INFERENCE_SERVER_BIN_DIR}" || error_exit "Failed to create bin directory"

	# Clone the repository
	log "Cloning tt-inference-server repository..."
	if [[ -d "${INFERENCE_SERVER_LIB_DIR}/tt-inference-server" ]]; then
		warn "tt-inference-server directory already exists at ${INFERENCE_SERVER_LIB_DIR}/tt-inference-server"
		warn "Skipping clone, will create wrapper script only"
	else
		git clone "${INFERENCE_SERVER_REPO_URL}" "${INFERENCE_SERVER_LIB_DIR}/tt-inference-server" || error_exit "Failed to clone tt-inference-server"
	fi

	# Create wrapper script
	log "Creating wrapper script..."
	cat > "${INFERENCE_SERVER_BIN_DIR}/${INFERENCE_SERVER_SCRIPT_NAME}" << 'EOF'
#!/bin/bash

cd ${HOME}/.local/lib/tt-inference-server
python ${HOME}/.local/lib/tt-inference-server/run.py "$@"
EOF

	# Make the script executable
	chmod +x "${INFERENCE_SERVER_BIN_DIR}/${INFERENCE_SERVER_SCRIPT_NAME}" || error_exit "Failed to make script executable"

	# Check if the directory is in PATH
	if [[ ":${PATH}:" != *":${INFERENCE_SERVER_BIN_DIR}:"* ]]; then
		warn "${INFERENCE_SERVER_BIN_DIR} is not in your PATH."
		warn "A restart may fix this, or you may need to update your shell RC"
	fi

	log "tt-inference-server installation completed"
	return 0
}

install_studio () {
	log "Installing tt-studio"
	local STUDIO_LIB_DIR="${HOME}/.local/lib"
	local STUDIO_BIN_DIR="${HOME}/.local/bin"
	local STUDIO_SCRIPT_NAME="tt-studio"
	local STUDIO_REPO_URL="https://github.com/tenstorrent/tt-studio.git"

	# Create directories
	mkdir -p "${STUDIO_LIB_DIR}" || error_exit "Failed to create library directory"
	mkdir -p "${STUDIO_BIN_DIR}" || error_exit "Failed to create bin directory"

	# Clone the repository
	log "Cloning tt-studio repository..."
	if [[ -d "${STUDIO_LIB_DIR}/tt-studio" ]]; then
		warn "tt-studio directory already exists at ${STUDIO_LIB_DIR}/tt-studio"
		warn "Skipping clone, will create wrapper script only"
	else
		git clone "${STUDIO_REPO_URL}" "${STUDIO_LIB_DIR}/tt-studio" || error_exit "Failed to clone tt-studio"
	fi

	# Create wrapper script
	log "Creating wrapper script..."
	cat > "${STUDIO_BIN_DIR}/${STUDIO_SCRIPT_NAME}" << 'EOF'
#!/bin/bash

cd "${HOME}/.local/lib/tt-studio"
python "${HOME}/.local/lib/tt-studio/run.py" "$@"
EOF

	# Make the script executable
	chmod +x "${STUDIO_BIN_DIR}/${STUDIO_SCRIPT_NAME}" || error_exit "Failed to make script executable"

	# Check if the directory is in PATH
	if [[ ":${PATH}:" != *":${STUDIO_BIN_DIR}:"* ]]; then
		warn "${STUDIO_BIN_DIR} is not in your PATH."
		warn "A restart may fix this, or you may need to update your shell RC"
	fi

	log "tt-studio installation completed"
	return 0
}

# Main installation script
main() {
	echo -e "${LOGO}"
	echo # newline
	INSTALLER_VERSION="__INSTALLER_DEVELOPMENT_BUILD__" # Set to semver at release time by GitHub Actions
	log "Welcome to tenstorrent!"
	log "This is tt-installer version ${INSTALLER_VERSION}"
	log "Log is at ${LOG_FILE}"

	log "This script will install drivers and tooling and properly configure your tenstorrent hardware."

	maybe_enable_default_mode
	log "Starting installation"

	# Log special mode settings
	if [[ "${_arg_mode_non_interactive}" = "on" ]]; then
		warn "Running in non-interactive mode"
	fi
	if [[ "${_arg_mode_container}" = "on" ]]; then
		warn "Running in container mode"
	fi
	if [[ "${_arg_install_kmd}" = "off" ]]; then
		warn "KMD installation will be skipped"
	fi
	if [[ "${_arg_install_hugepages}" = "off" ]]; then
		warn "HugePages setup will be skipped"
	fi
	if [[ "${_arg_install_podman}" = "off" ]]; then
		warn "Podman installation will be skipped"
	fi
	if [[ "${_arg_install_metalium_container}" = "off" ]]; then
		warn "Metalium installation will be skipped"
	fi
	if [[ "${_arg_install_sfpi}" = "off" ]]; then
		warn "SFPI installation will be skipped"
	fi
	if [[ "${_arg_install_inference_server}" = "off" ]]; then
		warn "tt-inference-server installation will be skipped"
	fi
	if [[ "${_arg_install_studio}" = "off" ]]; then
		warn "tt-studio installation will be skipped"
	fi
	# shellcheck disable=SC2154
	if [[ "${_arg_install_tt_flash}" = "off" ]]; then
		warn "TT-Flash installation will be skipped"
	fi
	if [[ "${_arg_update_firmware}" = "off" ]]; then
		warn "Firmware update will be skipped"
	fi
	if [[ "${_arg_update_firmware}" = "force" ]]; then
		warn "Firmware will be forcibly updated"
	fi
	if [[ "${_arg_install_metalium_models_container}" = "on" ]]; then
		log "Metalium Models container will be installed"
	fi
	if [[ "${_arg_use_uv}" = "on" ]]; then
		log "uv will be used instead of pip for package installation"
	fi

	log "Checking for sudo permissions... (may request password)"
	check_has_sudo_perms

	# Check distribution and install base packages
	detect_distro

	log "Installing base packages"
	case "${DISTRO_ID}" in
		"ubuntu")
			sudo apt-get update
			sudo DEBIAN_FRONTEND=noninteractive apt-get install -y git python3-pip dkms cargo rustc pipx jq protobuf-compiler wget
			;;
		"debian")
			# On Debian, packaged cargo and rustc are very old. Users must install them another way.
			sudo apt-get update
			sudo apt-get install -y git python3-pip dkms pipx jq protobuf-compiler wget
			;;
		"fedora")
			sudo dnf install -y git python3-pip python3-devel dkms cargo rust pipx jq protobuf-compiler wget
			;;
		"rhel"|"centos")
			sudo dnf install -y epel-release
			sudo dnf install -y git python3-pip python3-devel dkms cargo rust pipx jq protobuf-compiler wget
			;;
		*)
			error "Unsupported distribution: ${DISTRO_ID}"
			exit 1
			;;
	esac

	if [[ "${DISTRO_ID}" = "debian" ]]; then
		warn "rustc and cargo cannot be automatically installed on Debian. Ensure the latest versions are installed before continuing."
		warn "If you are unsure how to do this, use rustup: https://rustup.rs/"
	fi

	# Get Podman Metalium installation choice
	get_podman_metalium_choice

	# Get tt-inference-server installation choice
	get_inference_server_choice
	get_studio_choice

	# Python package installation preference
	get_python_choice
	install_tt_repos

	# 1. Define the package registry
	# Format: "package_name|install_flag|version|type"
	declare -A package_registry=(
		# System packages
		["kmd"]="tenstorrent-dkms|${_arg_install_kmd}|${_arg_kmd_version}|system"
		["hugepages"]="tenstorrent-tools|${_arg_install_hugepages}|${_arg_systools_version}|system"
		["sfpi"]="sfpi|${_arg_install_sfpi}|${_arg_sfpi_version}|system"
		["podman"]="podman|${_arg_install_podman}||system"
		["podman-docker"]="podman-docker|${_arg_install_podman_docker}||system"
		["podman-compose"]="podman-compose|${_arg_install_podman_docker}||system"

		# Python packages
		["tt-topology"]="tt-topology|${_arg_install_tt_topology}|${_arg_topology_version}|python"
		["tt-flash"]="tt-flash|${_arg_install_tt_flash}|${_arg_flash_version}|python"
		["tt-smi"]="tt-smi|${_arg_install_tt_smi}|${_arg_smi_version}|python"
	)

	# 2. Parse the registry to obtain lists of packages
	declare -a system_packages=()
	declare -a python_packages=()

	for key in "${!package_registry[@]}"; do
		IFS='|' read -r pkg_name install_flag version pkg_type <<< "${package_registry[${key}]}"

		# Skip if not marked for installation
		[[ "${install_flag}" != "on" ]] && continue

		# Add to appropriate list with version formatting
		case "${pkg_type}" in
			system)
				if [[ -z "${version}" ]]; then
					system_packages+=("${pkg_name}")
				else
					# Format based on package manager
					if [[ "${PKG_MANAGER}" = "apt-get" ]]; then
						system_packages+=("${pkg_name}=${version}")
					elif [[ "${PKG_MANAGER}" = "dnf" ]]; then
						system_packages+=("${pkg_name}-${version}")
					else
						system_packages+=("${pkg_name}")  # fallback to no version
					fi
				fi
				;;
			python)
				if [[ -z "${version}" ]]; then
					python_packages+=("${pkg_name}")
				else
					python_packages+=("${pkg_name}==${version}")
				fi
				;;
		esac
	done

	# 3. Act on the lists
	# Install system packages
	if [[ ${#system_packages[@]} -gt 0 ]]; then
		echo "Installing system packages: ${system_packages[*]}"
		if [[ "${PKG_MANAGER}" = "apt-get" ]]; then
			sudo apt-get install -y "${system_packages[@]}"
		elif [[ "${PKG_MANAGER}" = "dnf" ]]; then
			sudo dnf install -y "${system_packages[@]}"
		fi
	fi

	# Install Python packages
	if [[ ${#python_packages[@]} -gt 0 ]]; then
		echo "Installing Python packages: ${python_packages[*]}"
		if [[ -z "${PYTHON_INSTALL_CMD:-}" ]]; then
			error_exit "PYTHON_INSTALL_CMD is not set. Python package installation cannot proceed."
		fi
		${PYTHON_INSTALL_CMD} "${python_packages[@]}"
	fi

	# Update firmware using tt-flash
	if [[ "${_arg_update_firmware}" = "off" ]]; then
		log "Skipping firmware update"
	else
		log "Updating firmware"

		# Check if tt-flash is installed and available
		if ! command -v tt-flash &> /dev/null; then
			error_exit "tt-flash is not installed or not in PATH. Please install tt-flash before attempting firmware update."
		fi

		if [[ -n "${_arg_fw_version:-}" ]]; then
			FW_VERSION=${_arg_fw_version}
		else
			FW_VERSION=$(fetch_latest_version "tenstorrent/tt-firmware" "v");
		fi

		cd "${WORKDIR}"

		# Create FW_FILE based on FW_VERSION
		FW_FILE="fw_pack-${FW_VERSION}.fwbundle"
		FW_RELEASE_URL="https://github.com/tenstorrent/tt-firmware/releases/download"
		BACKUP_FW_RELEASE_URL="https://github.com/tenstorrent/tt-zephyr-platforms/releases/download"

		# Download from GitHub releases
		if ! curl -fsSLO "${FW_RELEASE_URL}/v${FW_VERSION}/${FW_FILE}"; then
			warn "Tried URL ${FW_RELEASE_URL}/v${FW_VERSION}/${FW_FILE}"
			warn "Could not find firmware bundle at main URL- trying backup URL"
			if ! curl -fsSLO "${BACKUP_FW_RELEASE_URL}/v${FW_VERSION}/${FW_FILE}"; then
				error_exit "Could not download firmware bundle. Ensure firmware version is valid."
			fi
		fi

		verify_download "${FW_FILE}"

		if [[ "${_arg_update_firmware}" = "force" ]]; then
			tt-flash flash "${FW_FILE}" --force
		else
			tt-flash flash "${FW_FILE}"
		fi
	fi

	if [[ "${_arg_install_inference_server}" = "on" ]]; then
		install_inference_server
	fi

	if [[ "${_arg_install_studio}" = "on" ]]; then
		install_studio
	fi

	# Setup rootless Podman if it was just installed
	if [[ "${_arg_install_podman}" = "on" ]]; then
		if check_podman_installed; then
			setup_rootless_podman
		else
			warn "Podman was not installed successfully"
		fi
	fi

	# Install Podman Metalium if requested
	if [[ "${_arg_install_metalium_container}" = "off" ]]; then
		warn "Skipping Podman Metalium installation"
	else
		if ! check_podman_installed; then
			warn "Podman is not installed. Cannot install Podman Metalium."
		else
			install_podman_metalium
		fi
	fi

	# Install Metalium Models container if requested
	if [[ "${_arg_install_metalium_models_container}" = "on" ]]; then
		if ! check_podman_installed; then
			warn "Podman is not installed. Cannot install Metalium Models."
		else
			install_podman_metalium_models
		fi
	fi

	if [[ "${INSTALLED_IN_VENV}" = "0" ]]; then
		warn "You'll need to run \"source ${VIRTUAL_ENV}/bin/activate\" to use tenstorrent's Python tools."
	fi

	log "Please reboot your system to complete the setup."
	log "After rebooting, try running 'tt-smi' to see the status of your hardware."
	if [[ "${_arg_install_metalium_container}" = "on" ]]; then
		log "Use 'tt-metalium' to access the Metalium programming environment"
		log "Usage examples:"
		log "  tt-metalium                   # Start an interactive shell"
		log "  tt-metalium [command]         # Run a specific command"
		log "  tt-metalium python script.py  # Run a Python script"
	fi
	if [[ "${_arg_install_inference_server}" = "on" ]]; then
		log "Use 'tt-inference-server' to run the inference server"
		log "The inference server has been installed to ~/.local/lib/tt-inference-server"
		log "Usage: tt-inference-server [arguments]"
	fi
	if [[ "${_arg_install_studio}" = "on" ]]; then
		log "Use 'tt-studio' to launch tt-studio"
		log "tt-studio has been installed to ~/.local/lib/tt-studio"
		log "Usage: tt-studio [arguments]"
	fi

	# Log successful completion message
	log "✅ Installation completed successfully."
	log "Installation log saved to: ${LOG_FILE}"

	# Auto-reboot if specified
	if [[ "${_arg_reboot_option}" = "always" ]]; then
		log "Auto-reboot enabled. Rebooting now..."
		sudo reboot
	# Otherwise, ask if specified
	elif [[ "${_arg_reboot_option}" = "ask" ]]; then
		if confirm "Would you like to reboot now?"; then
			log "Rebooting..."
			sudo reboot
		fi
	fi
}

# Start installation
main

# ] <-- needed because of Argbash

# vim: noai:ts=4:sw=4:ft=bash