#!/bin/bash

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

# Check if running as root
if [[ $EUID -ne 0 ]]; then
   log_error "This script must be run as root (use sudo)"
fi

# 1. Prerequisites
log_info "Updating package index and installing prerequisites..."
apt-get update
apt-get install -y ca-certificates curl gnupg lsb-release

# 2. Add Docker's official GPG key
log_info "Adding Docker's official GPG key..."
install -m 0755 -d /etc/apt/keyrings
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

# 3. Set up the repository
log_info "Setting up the Docker repository..."
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null

# 4. Install Docker Engine and CLI
log_info "Installing Docker Engine, CLI, and Containerd..."
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

# Note: Docker Compose V2 is now installed as a plugin (docker compose), 
# not the standalone binary (docker-compose).

# 5. Enable and Start Docker Service
log_info "Enabling and starting Docker service..."
systemctl enable docker.service
systemctl enable containerd.service
systemctl start docker

# 6. Add current user to docker group
# We detect the user running the script (sudo user) or the first non-root user
CURRENT_USER="${SUDO_USER:-$(whoami)}"
if id -nG "$CURRENT_USER" | grep -qw "docker"; then
    log_info "User '$CURRENT_USER' is already in the docker group."
else
    log_info "Adding user '$CURRENT_USER' to the docker group..."
    usermod -aG docker "$CURRENT_USER"
    log_warn "You must log out and log back in (or run 'newgrp docker') for group changes to take effect."
fi

# 7. Verification
log_info "Verifying installation..."
if docker --version >/dev/null 2>&1; then
    log_info "Docker installed successfully: $(docker --version)"
else
    log_error "Docker installation verification failed."
fi

if docker compose version >/dev/null 2>&1; then
    log_info "Docker Compose Plugin installed successfully: $(docker compose version)"
else
    log_error "Docker Compose installation verification failed."
fi

log_info "Installation complete!"
log_warn "Remember to run 'sudo usermod -aG docker $CURRENT_USER' if you didn't run as the correct user, and then log out/in."