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

### CONFIG ###
OPENWEBUI_IMAGE="ghcr.io/open-webui/open-webui:main"
OLLAMA_MODELS=("llama3.1:8b" "codellama:7b" "qwen2-math:7b")
### /CONFIG ###

echo "=== Step 2: Install Docker, Ollama, Stable Diffusion, and Open WebUI on Arch Linux ==="

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

if ! command -v pacman >/dev/null 2>&1; then
  echo "This script is for Arch Linux. pacman was not found."
  exit 1
fi

# Figure out the real user so user-owned app folders are created in their home directory.
REALUSER="${SUDO_USER:-$USER}"
USER_HOME="$(getent passwd "${REALUSER}" | cut -d: -f6)"
OPENWEBUI_DIR="${USER_HOME}/open-webui"
STABLEDIFF_DIR="${USER_HOME}/stablediff"

echo "Real user:       ${REALUSER}"
echo "User home:       ${USER_HOME}"
echo "Open WebUI dir:  ${OPENWEBUI_DIR}"
echo "Stable Diff dir: ${STABLEDIFF_DIR}"
echo

# --- Sanity check: GPU should be visible now ---
if ! command -v nvidia-smi >/dev/null 2>&1; then
  echo "WARNING: nvidia-smi not found."
  echo "Make sure you ran the Arch NVIDIA driver script and rebooted."
  echo "You can continue, but GPU usage may not work."
else
  echo "nvidia-smi output:"
  nvidia-smi || true
fi

###############################################################################
# Step 1/4: Install Docker Engine + Docker Compose
###############################################################################
echo
echo "Step 1/4: Installing Docker Engine + Docker Compose..."

pacman -Syu --noconfirm
pacman -S --needed --noconfirm \
  docker \
  docker-compose \
  docker-buildx \
  containerd \
  ocl-icd \
  opencl-headers

systemctl enable --now docker

# Allow the real user to run docker without sudo after logging out/in.
if id -nG "${REALUSER}" | grep -qw docker; then
  echo "${REALUSER} is already in the docker group."
else
  usermod -aG docker "${REALUSER}"
  echo "Added ${REALUSER} to docker group. Log out/in for this to fully apply."
fi

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

###############################################################################
# Step 2/4: Install Ollama
###############################################################################
echo
echo "Step 2/4: Installing Ollama..."

# Prefer the CUDA-enabled Arch package when available; fall back to plain ollama.
if pacman -Si ollama-cuda >/dev/null 2>&1; then
  pacman -S --needed --noconfirm ollama-cuda
else
  pacman -S --needed --noconfirm ollama
fi

systemctl enable --now ollama || true

echo "Waiting for Ollama API 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"
fi

###############################################################################
# Step 3/4: Install Stable Diffusion AUTOMATIC1111 via pyenv + systemd autostart
###############################################################################
echo
echo "Step 3/4: Installing Stable Diffusion WebUI AUTOMATIC1111 with pyenv..."

echo "Installing prerequisites for Stable Diffusion..."
pacman -S --needed --noconfirm \
  base-devel \
  openssl \
  zlib \
  bzip2 \
  readline \
  sqlite \
  wget \
  curl \
  llvm \
  ncurses \
  xz \
  tk \
  libffi \
  git \
  python-pip \
  pyenv

# Configure pyenv for the REALUSER if not already present.
sudo -u "${REALUSER}" bash -lc '
  grep -q "PYENV_ROOT" "$HOME/.bashrc" 2>/dev/null || cat >> "$HOME/.bashrc" <<"BASHRC_EOF"

# pyenv
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"
BASHRC_EOF
'

# Use pyenv to ensure Python 3.10 for Stable Diffusion.
sudo -u "${REALUSER}" bash -lc '
  export PYENV_ROOT="$HOME/.pyenv"
  export PATH="$PYENV_ROOT/bin:/usr/bin:$PATH"
  eval "$(pyenv init -)"

  echo "Ensuring Python 3.10 is installed with pyenv..."
  pyenv install -s 3.10
  pyenv global 3.10
'

# Clone AUTOMATIC1111 repo into ~/stablediff if needed.
sudo -u "${REALUSER}" bash -lc '
  if [ ! -d "$HOME/stablediff" ]; then
    echo "Cloning Stable Diffusion WebUI repo into ~/stablediff..."
    git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git "$HOME/stablediff"
  else
    echo "~/stablediff already exists, skipping clone."
  fi
'

# Create systemd service for Stable Diffusion WebUI to auto-start at boot.
echo "Creating systemd service for Stable Diffusion WebUI..."

cat >/etc/systemd/system/stablediff.service <<EOF_SERVICE
[Unit]
Description=Stable Diffusion WebUI AUTOMATIC1111
After=network-online.target nvidia-persistenced.service
Wants=network-online.target

[Service]
Type=simple
User=${REALUSER}
WorkingDirectory=${STABLEDIFF_DIR}
Environment=PYENV_ROOT=${USER_HOME}/.pyenv
Environment=PATH=${USER_HOME}/.pyenv/shims:${USER_HOME}/.pyenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
Environment=PYTHON=${USER_HOME}/.pyenv/shims/python
ExecStart=/bin/bash -lc './webui.sh --listen --api'
Restart=on-failure
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF_SERVICE

systemctl daemon-reload
systemctl enable --now stablediff.service || true

echo "Stable Diffusion WebUI service created: stablediff.service"
echo "Check status with: systemctl status stablediff.service"

###############################################################################
# Step 4/4: Set up Open WebUI with Docker Compose talking to Ollama
###############################################################################
echo
echo "Step 4/4: Setting up Open WebUI with Docker Compose..."

mkdir -p "${OPENWEBUI_DIR}"
chown -R "${REALUSER}:${REALUSER}" "${OPENWEBUI_DIR}"
cd "${OPENWEBUI_DIR}"

cat > docker-compose.yml <<EOF_COMPOSE
services:
  open-webui:
    image: ${OPENWEBUI_IMAGE}
    container_name: open-webui
    restart: unless-stopped
    network_mode: "host"
    environment:
      - OLLAMA_BASE_URL=http://127.0.0.1:11434
      # Uncomment to disable auth on local networks only. Do not expose this publicly with auth disabled.
      # - OPENWEBUI_AUTH=False
    volumes:
      - ./data:/app/backend/data
EOF_COMPOSE

chown "${REALUSER}:${REALUSER}" docker-compose.yml

docker compose up -d

###############################################################################
# Pull starter Ollama models
###############################################################################
echo
echo "Pulling starter Ollama models..."
for model in "${OLLAMA_MODELS[@]}"; do
  ollama pull "$model" || echo "WARNING: failed to pull $model"
done

###############################################################################
# Done
###############################################################################
echo
echo "==================== ALL DONE ===================="
echo "Ollama service:              systemctl status ollama"
echo "Stable Diffusion WebUI:      systemctl status stablediff.service"
echo "Open WebUI stack:            cd ${OPENWEBUI_DIR} && docker compose ps"
echo
echo "Web UIs:"
echo "  Open WebUI via Ollama:          http://<your-server-ip>:8080"
echo "  Stable Diffusion WebUI:         http://<your-server-ip>:7860"
echo
echo "Note: if Docker permissions do not work for ${REALUSER}, log out and back in."
echo "==================================================="