#!/bin/bash
#
# Docker Installation Script for Raspberry Pi OS 64-bit Lite
# Usage: sudo bash install_docker_rpi.sh
#

set -euo pipefail

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

if [ "$(id -u)" -ne 0 ]; then
    log "ERROR: This script must be run as root (use sudo)."
    exit 1
fi

SWARM_INIT=false
SWARM_JOIN=false
SWARM_MANAGER_IP=""
SWARM_TOKEN=""

echo "=== Docker Swarm Configuration ==="
read -r -p "Is this the first node of a new Docker Swarm? (y/N): " init_choice
if [[ "$init_choice" =~ ^[Yy]$ ]]; then
    SWARM_INIT=true
else
    read -r -p "Do you want to join an existing Docker Swarm? (y/N): " join_choice
    if [[ "$join_choice" =~ ^[Yy]$ ]]; then
        SWARM_JOIN=true
        read -r -p "Enter Swarm Manager IP address: " SWARM_MANAGER_IP
        read -r -p "Enter Swarm Join Token: " SWARM_TOKEN
    fi
fi

log "Updating system packages..."
apt-get update
apt-get upgrade -y

log "Downloading and running Docker's official convenience script..."
curl -fsSL https://get.docker.com -o /tmp/get-docker.sh
sh /tmp/get-docker.sh

log "Configuring user permissions..."
# Use the user who invoked sudo, otherwise default to 'pi'
TARGET_USER="${SUDO_USER:-pi}"

if id "$TARGET_USER" &>/dev/null; then
    usermod -aG docker "$TARGET_USER"
    log "Added user '$TARGET_USER' to the 'docker' group."
else
    log "WARNING: User '$TARGET_USER' not found, skipping group assignment."
fi

log "Enabling Docker service to start on boot..."
systemctl enable docker
systemctl start docker

# Grab the primary local IP address to advertise for Swarm
LOCAL_IP=$(hostname -I | awk '{print $1}')

if [ "$SWARM_INIT" = true ]; then
    log "Initializing Docker Swarm..."
    docker swarm init --advertise-addr "$LOCAL_IP"
    WORKER_TOKEN=$(docker swarm join-token worker -q)
elif [ "$SWARM_JOIN" = true ]; then
    log "Joining Docker Swarm at $SWARM_MANAGER_IP..."
    # Append default swarm port if the user didn't specify one
    if [[ "$SWARM_MANAGER_IP" != *":"* ]]; then
        SWARM_MANAGER_IP="${SWARM_MANAGER_IP}:2377"
    fi
    docker swarm join --advertise-addr "$LOCAL_IP" --token "$SWARM_TOKEN" "$SWARM_MANAGER_IP"
fi

log "Cleaning up..."
rm -f /tmp/get-docker.sh

log "=== Installation Complete ==="
log "Note: You must log out and log back in for the group changes to take effect,"
log "      or run 'newgrp docker' to use Docker as a non-root user immediately."

if [ "$SWARM_INIT" = true ]; then
    echo ""
    log "=== Swarm Initialized ==="
    log "To add worker nodes to this swarm, use the following as input variables:"
    log "Manager IP: $LOCAL_IP"
    log "Join Token: $WORKER_TOKEN"
fi