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

### CONFIG ###
OPENWEBUI_IMAGE="ghcr.io/open-webui/open-webui:main"
### /CONFIG ###

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

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

# Figure out the "real" user so we install things into their home directory
REALUSER="${SUDO_USER:-$USER}"
USER_HOME="$(eval echo ~${REALUSER})"
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 00-gpu-driver.sh and rebooted successfully."
  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..."

# Remove old Docker bits if present (ignore errors)
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 \
  ocl-icd-opencl-dev \
  opencl-headers

systemctl enable --now docker

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

###############################################################################
# Step 2/4: Install Ollama (GPU-enabled)
###############################################################################
echo
echo "Step 2/4: Installing Ollama (GPU-enabled)..."

# Let Ollama's installer run even if its CUDA driver step conflicts with Ubuntu's drivers
set +e
curl -fsSL https://ollama.com/install.sh | sh
OLLAMA_RC=$?
set -e

if [ $OLLAMA_RC -ne 0 ]; then
  echo "WARNING: Ollama installer returned non-zero (likely CUDA driver dependency issues)."
  echo "As long as your NVIDIA driver works (nvidia-smi OK), you can ignore this."
fi

# Make sure the service is enabled & running
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..."
apt-get install -y \
  make build-essential libssl-dev zlib1g-dev libbz2-dev \
  libreadline-dev libsqlite3-dev wget curl llvm libncurses5-dev \
  libncursesw5-dev xz-utils tk-dev libffi-dev liblzma-dev git \
  python3-pip

# Install pyenv for the REALUSER if not already present
sudo -u "${REALUSER}" bash -lc '
  if [ ! -d "$HOME/.pyenv" ]; then
    echo "Installing pyenv for $USER..."
    curl https://pyenv.run | bash
  else
    echo "pyenv already installed for $USER, skipping."
  fi
'

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

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

# Clone AUTOMATIC1111’s repo into ~/stablediff (if not already there)
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
[Unit]
Description=Stable Diffusion WebUI (AUTOMATIC1111)
After=network.target nvidia-persistenced.service

[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

[Install]
WantedBy=multi-user.target
EOF

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}"
cd "${OPENWEBUI_DIR}"

cat > docker-compose.yml <<'EOF'
version: "3.8"

services:
  open-webui:
    image: ghcr.io/open-webui/open-webui:main
    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 (optional, not recommended on exposed hosts)
      # - OPENWEBUI_AUTH=False
    volumes:
      - ./data:/app/backend/data
EOF

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

docker compose up -d

###############################################################################
# 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 (LLMs via Ollama):    http://<your-server-ip>:8080"
echo "  Stable Diffusion WebUI (images): http://<your-server-ip>:7860"
echo
echo "Stable Diffusion is configured to auto-start on boot via:"
echo "  systemctl enable stablediff.service"
echo "==================================================="

ollama pull llama3.1:8b
ollama pull codellama:7b
ollama pull qwen2-math:7b