#!/bin/bash
set -euo pipefail

sudo apt update
sudo apt upgrade -y
sudo apt install -y linux-headers-$(uname -r) build-essential dkms

sudo apt purge -y 'cuda-drivers*' 'nvidia-driver-*server*' || true

sudo ubuntu-drivers install --gpgpu
sudo apt install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker

# Ollama
curl -fsSL https://ollama.com/install.sh | sh

# Docker Engine + Compose
sudo apt-get update
sudo apt-get install -y ca-certificates curl
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo tee /etc/apt/keyrings/docker.asc >/dev/null
sudo chmod a+r /etc/apt/keyrings/docker.asc
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] \
https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" \
| sudo tee /etc/apt/sources.list.d/docker.list >/dev/null
sudo apt-get update
sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

# Run Docker without sudo from now on
sudo usermod -aG docker $USER && newgrp docker

# Option B: compose file
mkdir -p ~/docker-compose-files/open_web_ui && cd ~/docker-compose-files/open_web_ui
wget https://wget-pulls.fortis-scientia.com/open_web_ui.yaml
sudo docker compose -f open_web_ui.yaml up -d

# Pull models (note: llama3.3:70b is huge; ensure VRAM/disk)
ollama pull llama3:8b
ollama pull codellama:7b
ollama pull qwen2-math:7b