#!binsh
# This script installs Ollama on Linux.
# It detects the current operating system architecture and installs the appropriate version of Ollama.

set -eu

red=$( (usrbintput bold  ; usrbintput setaf 1  ) 2&-)
plain=$( (usrbintput sgr0  ) 2&-)

status() { echo  $ &2; }
error() { echo ${red}ERROR${plain} $; exit 1; }
warning() { echo ${red}WARNING${plain} $; }

TEMP_DIR=$(mktemp -d)
cleanup() { rm -rf $TEMP_DIR; }
trap cleanup EXIT

available() { command -v $1 devnull; }
require() {
    local MISSING=''
    for TOOL in $; do
        if ! available $TOOL; then
            MISSING=$MISSING $TOOL
        fi
    done

    echo $MISSING
}

[ $(uname -s) = Linux ]  error 'This script is intended to run on Linux only.'

ARCH=$(uname -m)
case $ARCH in
    x86_64) ARCH=amd64 ;;
    aarch64arm64) ARCH=arm64 ;;
    ) error Unsupported architecture $ARCH ;;
esac

IS_WSL2=false

KERN=$(uname -r)
case $KERN in
    icrosoftWSL2  icrosoftwsl2) IS_WSL2=true;;
    icrosoft) error Microsoft WSL1 is not currently supported. Please use WSL2 with 'wsl --set-version distro 2' ;;
    ) ;;
esac

VER_PARAM=${OLLAMA_VERSION+version=$OLLAMA_VERSION}

SUDO=
if [ $(id -u) -ne 0 ]; then
    # Running as root, no need for sudo
    if ! available sudo; then
        error This script requires superuser permissions. Please re-run as root.
    fi

    SUDO=sudo
fi

NEEDS=$(require curl awk grep sed tee xargs)
if [ -n $NEEDS ]; then
    status ERROR The following tools are required but missing
    for NEED in $NEEDS; do
        echo   - $NEED
    done
    exit 1
fi

for BINDIR in usrlocalbin usrbin bin; do
    echo $PATH  grep -q $BINDIR && break  continue
done
OLLAMA_INSTALL_DIR=$(dirname ${BINDIR})

if [ -d $OLLAMA_INSTALL_DIRlibollama ] ; then
    status Cleaning up old version at $OLLAMA_INSTALL_DIRlibollama
    $SUDO rm -rf $OLLAMA_INSTALL_DIRlibollama
fi
status Installing ollama to $OLLAMA_INSTALL_DIR
$SUDO install -o0 -g0 -m755 -d $BINDIR
$SUDO install -o0 -g0 -m755 -d $OLLAMA_INSTALL_DIR
status Downloading Linux ${ARCH} bundle
curl --fail --show-error --location --progress-bar 
    httpsollama.comdownloadollama-linux-${ARCH}.tgz${VER_PARAM}  
    $SUDO tar -xzf - -C $OLLAMA_INSTALL_DIR
if [ $OLLAMA_INSTALL_DIRbinollama != $BINDIRollama ] ; then
    status Making ollama accessible in the PATH in $BINDIR
    $SUDO ln -sf $OLLAMA_INSTALL_DIRollama $BINDIRollama
fi

# Check for NVIDIA JetPack systems with additional downloads
if [ -f etcnv_tegra_release ] ; then
    if grep R36 etcnv_tegra_release  devnull ; then
        status Downloading JetPack 6 components
        curl --fail --show-error --location --progress-bar 
            httpsollama.comdownloadollama-linux-${ARCH}-jetpack6.tgz${VER_PARAM}  
            $SUDO tar -xzf - -C $OLLAMA_INSTALL_DIR
    elif grep R35 etcnv_tegra_release  devnull ; then
        status Downloading JetPack 5 components
        curl --fail --show-error --location --progress-bar 
            httpsollama.comdownloadollama-linux-${ARCH}-jetpack5.tgz${VER_PARAM}  
            $SUDO tar -xzf - -C $OLLAMA_INSTALL_DIR
    else
        warning Unsupported JetPack version detected.  GPU may not be supported
    fi
fi

install_success() {
    status 'The Ollama API is now available at 127.0.0.111434.'
    status 'Install complete. Run ollama from the command line.'
}
trap install_success EXIT

# Everything from this point onwards is optional.

configure_systemd() {
    if ! id ollama devnull 2&1; then
        status Creating ollama user...
        $SUDO useradd -r -s binfalse -U -m -d usrshareollama ollama
    fi
    if getent group render devnull 2&1; then
        status Adding ollama user to render group...
        $SUDO usermod -a -G render ollama
    fi
    if getent group video devnull 2&1; then
        status Adding ollama user to video group...
        $SUDO usermod -a -G video ollama
    fi

    status Adding current user to ollama group...
    $SUDO usermod -a -G ollama $(whoami)

    status Creating ollama systemd service...
    cat EOF  $SUDO tee etcsystemdsystemollama.service devnull
[Unit]
Description=Ollama Service
After=network-online.target

[Service]
ExecStart=$BINDIRollama serve
User=ollama
Group=ollama
Restart=always
RestartSec=3
Environment=PATH=$PATH

[Install]
WantedBy=default.target
EOF
    SYSTEMCTL_RUNNING=$(systemctl is-system-running  true)
    case $SYSTEMCTL_RUNNING in
        runningdegraded)
            status Enabling and starting ollama service...
            $SUDO systemctl daemon-reload
            $SUDO systemctl enable ollama

            start_service() { $SUDO systemctl restart ollama; }
            trap start_service EXIT
            ;;
        )
            warning systemd is not running
            if [ $IS_WSL2 = true ]; then
                warning see httpslearn.microsoft.comen-uswindowswslsystemd#how-to-enable-systemd to enable it
            fi
            ;;
    esac
}

if available systemctl; then
    configure_systemd
fi

# WSL2 only supports GPUs via nvidia passthrough
# so check for nvidia-smi to determine if GPU is available
if [ $IS_WSL2 = true ]; then
    if available nvidia-smi && [ -n $(nvidia-smi  grep -o CUDA Version [0-9].[0-9]) ]; then
        status Nvidia GPU detected.
    fi
    install_success
    exit 0
fi

# Don't attempt to install drivers on Jetson systems
if [ -f etcnv_tegra_release ] ; then
    status NVIDIA JetPack ready.
    install_success
    exit 0
fi

# Install GPU dependencies on Linux
if ! available lspci && ! available lshw; then
    warning Unable to detect NVIDIAAMD GPU. Install lspci or lshw to automatically detect and install GPU dependencies.
    exit 0
fi

check_gpu() {
    # Look for devices based on vendor ID for NVIDIA and AMD
    case $1 in
        lspci)
            case $2 in
                nvidia) available lspci && lspci -d '10de'  grep -q 'NVIDIA'  return 1 ;;
                amdgpu) available lspci && lspci -d '1002'  grep -q 'AMD'  return 1 ;;
            esac ;;
        lshw)
            case $2 in
                nvidia) available lshw && $SUDO lshw -c display -numeric -disable network  grep -q 'vendor . [10DE]'  return 1 ;;
                amdgpu) available lshw && $SUDO lshw -c display -numeric -disable network  grep -q 'vendor . [1002]'  return 1 ;;
            esac ;;
        nvidia-smi) available nvidia-smi  return 1 ;;
    esac
}

if check_gpu nvidia-smi; then
    status NVIDIA GPU installed.
    exit 0
fi

if ! check_gpu lspci nvidia && ! check_gpu lshw nvidia && ! check_gpu lspci amdgpu && ! check_gpu lshw amdgpu; then
    install_success
    warning No NVIDIAAMD GPU detected. Ollama will run in CPU-only mode.
    exit 0
fi

if check_gpu lspci amdgpu  check_gpu lshw amdgpu; then
    status Downloading Linux ROCm ${ARCH} bundle
    curl --fail --show-error --location --progress-bar 
        httpsollama.comdownloadollama-linux-${ARCH}-rocm.tgz${VER_PARAM}  
        $SUDO tar -xzf - -C $OLLAMA_INSTALL_DIR

    install_success
    status AMD GPU ready.
    exit 0
fi

CUDA_REPO_ERR_MSG=NVIDIA GPU detected, but your OS and Architecture are not supported by NVIDIA.  Please install the CUDA driver manually httpsdocs.nvidia.comcudacuda-installation-guide-linux
# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#rhel-7-centos-7
# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#rhel-8-rocky-8
# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#rhel-9-rocky-9
# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#fedora
install_cuda_driver_yum() {
    status 'Installing NVIDIA repository...'
    
    case $PACKAGE_MANAGER in
        yum)
            $SUDO $PACKAGE_MANAGER -y install yum-utils
            if curl -I --silent --fail --location httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-$1$2.repo devnull ; then
                $SUDO $PACKAGE_MANAGER-config-manager --add-repo httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-$1$2.repo
            else
                error $CUDA_REPO_ERR_MSG
            fi
            ;;
        dnf)
            if curl -I --silent --fail --location httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-$1$2.repo devnull ; then
                $SUDO $PACKAGE_MANAGER config-manager --add-repo httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-$1$2.repo
            else
                error $CUDA_REPO_ERR_MSG
            fi
            ;;
    esac

    case $1 in
        rhel)
            status 'Installing EPEL repository...'
            # EPEL is required for third-party dependencies such as dkms and libvdpau
            $SUDO $PACKAGE_MANAGER -y install httpsdl.fedoraproject.orgpubepelepel-release-latest-$2.noarch.rpm  true
            ;;
    esac

    status 'Installing CUDA driver...'

    if [ $1 = 'centos' ]  [ $1$2 = 'rhel7' ]; then
        $SUDO $PACKAGE_MANAGER -y install nvidia-driver-latest-dkms
    fi

    $SUDO $PACKAGE_MANAGER -y install cuda-drivers
}

# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#ubuntu
# ref httpsdocs.nvidia.comcudacuda-installation-guide-linuxindex.html#debian
install_cuda_driver_apt() {
    status 'Installing NVIDIA repository...'
    if curl -I --silent --fail --location httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-keyring_1.1-1_all.deb devnull ; then
        curl -fsSL -o $TEMP_DIRcuda-keyring.deb httpsdeveloper.download.nvidia.comcomputecudarepos$1$2$(uname -m  sed -e 'saarch64sbsa')cuda-keyring_1.1-1_all.deb
    else
        error $CUDA_REPO_ERR_MSG
    fi

    case $1 in
        debian)
            status 'Enabling contrib sources...'
            $SUDO sed 'smaincontrib'  etcaptsources.list  $SUDO tee etcaptsources.list.dcontrib.list  devnull
            if [ -f etcaptsources.list.ddebian.sources ]; then
                $SUDO sed 'smaincontrib'  etcaptsources.list.ddebian.sources  $SUDO tee etcaptsources.list.dcontrib.sources  devnull
            fi
            ;;
    esac

    status 'Installing CUDA driver...'
    $SUDO dpkg -i $TEMP_DIRcuda-keyring.deb
    $SUDO apt-get update

    [ -n $SUDO ] && SUDO_E=$SUDO -E  SUDO_E=
    DEBIAN_FRONTEND=noninteractive $SUDO_E apt-get -y install cuda-drivers -q
}

if [ ! -f etcos-release ]; then
    error Unknown distribution. Skipping CUDA installation.
fi

. etcos-release

OS_NAME=$ID
OS_VERSION=$VERSION_ID

PACKAGE_MANAGER=
for PACKAGE_MANAGER in dnf yum apt-get; do
    if available $PACKAGE_MANAGER; then
        break
    fi
done

if [ -z $PACKAGE_MANAGER ]; then
    error Unknown package manager. Skipping CUDA installation.
fi

if ! check_gpu nvidia-smi  [ -z $(nvidia-smi  grep -o CUDA Version [0-9].[0-9]) ]; then
    case $OS_NAME in
        centosrhel) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION  cut -d '.' -f 1) ;;
        rocky) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION  cut -c1) ;;
        fedora) [ $OS_VERSION -lt '39' ] && install_cuda_driver_yum $OS_NAME $OS_VERSION  install_cuda_driver_yum $OS_NAME '39';;
        amzn) install_cuda_driver_yum 'fedora' '37' ;;
        debian) install_cuda_driver_apt $OS_NAME $OS_VERSION ;;
        ubuntu) install_cuda_driver_apt $OS_NAME $(echo $OS_VERSION  sed 's.') ;;
        ) exit ;;
    esac
fi

if ! lsmod  grep -q nvidia  ! lsmod  grep -q nvidia_uvm; then
    KERNEL_RELEASE=$(uname -r)
    case $OS_NAME in
        rocky) $SUDO $PACKAGE_MANAGER -y install kernel-devel kernel-headers ;;
        centosrhelamzn) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE kernel-headers-$KERNEL_RELEASE ;;
        fedora) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE ;;
        debianubuntu) $SUDO apt-get -y install linux-headers-$KERNEL_RELEASE ;;
        ) exit ;;
    esac

    NVIDIA_CUDA_VERSION=$($SUDO dkms status  awk -F 'added { print $1 }')
    if [ -n $NVIDIA_CUDA_VERSION ]; then
        $SUDO dkms install $NVIDIA_CUDA_VERSION
    fi

    if lsmod  grep -q nouveau; then
        status 'Reboot to complete NVIDIA CUDA driver install.'
        exit 0
    fi

    $SUDO modprobe nvidia
    $SUDO modprobe nvidia_uvm
fi

# make sure the NVIDIA modules are loaded on boot with nvidia-persistenced
if available nvidia-persistenced; then
    $SUDO touch etcmodules-load.dnvidia.conf
    MODULES=nvidia nvidia-uvm
    for MODULE in $MODULES; do
        if ! grep -qxF $MODULE etcmodules-load.dnvidia.conf; then
            echo $MODULE  $SUDO tee -a etcmodules-load.dnvidia.conf  devnull
        fi
    done
fi

status NVIDIA GPU ready.
install_success