#!/bin/bash
set -euo pipefail

# Always run relative to this script's folder (so compose file is found)
SCRIPT_DIR="$(cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P)"
cd "$SCRIPT_DIR"

# --- System prep ---
sudo apt update
sudo apt -y upgrade

# Kernel build deps and headers meta (avoids uname mismatch after upgrade)
sudo apt install -y build-essential dkms linux-headers-generic

# --- NVIDIA driver (auto-select) ---
sudo ubuntu-drivers install --gpgpu

# Strongly recommended: reboot here so the new kernel/driver settle,
# then re-run this script to do the rest.
# exit 0
# OR uncomment to auto-reboot now (and rerun manually after boot):
# sudo reboot

# --- 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

# --- NVIDIA Container Toolkit (after Docker is installed) ---
sudo apt install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# Optional: verify driver is live (helps catch Secure Boot/MOK issues)
if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi >/dev/null 2>&1; then
  echo "Warning: NVIDIA driver not loaded yet (Secure Boot? reboot needed?). GPU in containers may not work."
fi

# --- Ollama (host install; runs on 127.0.0.1:11434) ---
if ! command -v ollama >/dev/null 2>&1; then
  curl -fsSL https://ollama.com/install.sh | sh
fi

# --- OpenWebUI via Compose (talks to host Ollama) ---
mkdir -p "$HOME/docker-compose-files/open_web_ui"

# Copy (don't move) the compose file that sits next to this script.
# -n avoids overwriting an existing file.
if [ -f "$SCRIPT_DIR/docker-compose.yaml" ]; then
  cp -n "$SCRIPT_DIR/docker-compose.yaml" "$HOME/docker-compose-files/open_web_ui/docker-compose.yaml"
elif [ -f "$SCRIPT_DIR/docker-compose.yml" ]; then
  cp -n "$SCRIPT_DIR/docker-compose.yml" "$HOME/docker-compose-files/open_web_ui/docker-compose.yml"
fi

cd "$HOME/docker-compose-files/open_web_ui"

# Validate compose before launch
sudo docker compose config >/dev/null

# Bring it up
sudo docker compose up -d

# --- Models (ensure Ollama service is running) ---
# (These will fail the script if a model name is wrong due to -e)
ollama pull codellama:7b
ollama pull qwen2-math:7b
ollama pull llama3:8b