#!/usr/bin/env bash
set -euo pipefail

### CONFIG ###
mkdir -p /home/anon/open-webui
OPENWEBUI_DIR="/home/anon/open-webui"
OPENWEBUI_IMAGE="ghcr.io/open-webui/open-webui:main"
### /CONFIG ###

echo "=== Ollama + Open WebUI installer for Ubuntu Server 24.04 (RTX 3060) ==="

if [[ $EUID -ne 0 ]]; then
  echo "Please run this script as root (e.g. sudo bash $0)"
  exit 1
fi

echo "Step 1/6: Updating apt and installing base dependencies..."
apt-get update -y
apt-get upgrade -y
apt-get install -y \
  curl \
  ca-certificates \
  gnupg \
  lsb-release \
  apt-transport-https \
  software-properties-common

echo "Step 2/6: Installing NVIDIA driver (for RTX 3060)..."
# Install the recommended proprietary driver
ubuntu-drivers install --gpgpu || true

if ! command -v nvidia-smi >/dev/null 2>&1; then
  echo "WARNING: nvidia-smi not found. A reboot may be required to load the NVIDIA driver."
  echo "After this script finishes, reboot your server, then run: nvidia-smi"
else
  echo "NVIDIA driver appears to be installed:"
  nvidia-smi || true
fi

echo "Step 3/6: Installing Docker Engine + Docker Compose..."

# Remove old Docker if present
apt-get remove -y docker docker-engine docker.io containerd runc || true

# Add Docker’s official GPG key
install -m 0755 -d /etc/apt/keyrings
if [ ! -f /etc/apt/keyrings/docker.gpg ]; then
  curl -fsSL https://download.docker.com/linux/ubuntu/gpg | \
    gpg --dearmor -o /etc/apt/keyrings/docker.gpg
  chmod a+r /etc/apt/keyrings/docker.gpg
fi

# Add Docker apt repository
UBUNTU_CODENAME="$(. /etc/os-release && echo "$VERSION_CODENAME")"
if [ ! -f /etc/apt/sources.list.d/docker.list ]; then
  echo \
    "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
    ${UBUNTU_CODENAME} stable" > /etc/apt/sources.list.d/docker.list
fi

apt-get update -y
apt-get install -y \
  docker-ce \
  docker-ce-cli \
  containerd.io \
  docker-buildx-plugin \
  docker-compose-plugin

systemctl enable --now docker

echo "Docker version:"
docker --version || true
echo "Docker Compose plugin version:"
docker compose version || true

echo "Step 4/6: Installing Ollama (GPU-enabled)..."
curl -fsSL https://ollama.com/install.sh | sh

# Make sure the service is enabled & running
systemctl enable --now ollama

echo "Waiting for Ollama API to be reachable on http://127.0.0.1:11434 ..."
for i in {1..30}; do
  if curl -sSf http://127.0.0.1:11434/api/version >/dev/null 2>&1; then
    echo "Ollama is up!"
    break
  fi
  sleep 2
done

if ! curl -sSf http://127.0.0.1:11434/api/version >/dev/null 2>&1; then
  echo "WARNING: Ollama API not responding yet. Check 'systemctl status ollama' after the script finishes."
fi

echo "Step 5/6: Creating Open WebUI Docker Compose setup..."

mkdir -p "${OPENWEBUI_DIR}"
cd "${OPENWEBUI_DIR}"

cat > docker-compose.yml <<'EOF'
services:
  open-webui:
    image: ghcr.io/open-webui/open-webui:main
    container_name: open-webui
    restart: unless-stopped
    network_mode: "host"
    environment:
      # Point Open WebUI to the locally-running Ollama instance
      - OLLAMA_BASE_URL=http://127.0.0.1:11434
      # Uncomment to disable auth on local networks (optional)
      # - OPENWEBUI_AUTH=False
    volumes:
      - ./data:/app/backend/data
EOF

echo "docker-compose.yml created at ${OPENWEBUI_DIR}/docker-compose.yml"

echo "Step 6/6: Starting Open WebUI with Docker Compose..."
docker compose up -d

echo
echo "==================== DONE ===================="
echo "Ollama service:   systemctl status ollama"
echo "Open WebUI stack: cd ${OPENWEBUI_DIR} && docker compose ps"
echo
echo "If this is a fresh NVIDIA driver install, reboot the server, then verify GPU availability with:"
echo "  nvidia-smi"
echo
echo "Once everything is running, access Open WebUI at:"
echo "  http://<your-server-ip>:8080"
echo "(Default Open WebUI port inside the container is 8080; we're using host networking.)"
echo "=============================================="