#!/bin/bash
set -e

# --- Colors and Formatting ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
RESET='\033[0m'

info()    { echo -e "${CYAN}ℹ  $*${RESET}"; }
success() { echo -e "${GREEN}✅ $*${RESET}"; }
warn()    { echo -e "${YELLOW}⚠  $*${RESET}"; }
error()   { echo -e "${RED}❌ $*${RESET}"; }
step()    { echo -e "${BOLD}${BLUE}➤ $*${RESET}"; }

# Prompt for sudo password upfront, so it caches for script duration
if ! sudo -v; then
    error "This script requires sudo privileges. Please run again with a sudo-capable user."
    exit 1
fi

# Keep-alive to update sudo timestamp until script finishes
while true; do sudo -n true; sleep 60; kill -0 "$$" || exit; done 2>/dev/null &

# --- Detect distribution ---
if [ -f /etc/os-release ]; then
    . /etc/os-release
    DISTRO_ID="$ID"
    DISTRO_VERSION="$VERSION_ID"
    DISTRO_CODENAME="$VERSION_CODENAME"
else
    error "Cannot detect Linux distribution."
    exit 1
fi

# Fallback for older Debian/Ubuntu releases
if [[ "$DISTRO_ID" == "debian" ]]; then
    if [[ -z "$DISTRO_CODENAME" || ! -f "/usr/share/keyrings/docker-archive-keyring.gpg" ]]; then
        DISTRO_CODENAME="bookworm"  # Fallback to Debian 12
    fi
elif [[ "$DISTRO_ID" == "ubuntu" ]]; then
    if [[ -z "$DISTRO_CODENAME" || ! -f "/usr/share/keyrings/docker-archive-keyring.gpg" ]]; then
        DISTRO_CODENAME="jammy"  # Fallback to Ubuntu 22.04 LTS
    fi
fi

install_docker() {
    step "Installing Docker for ${BOLD}$DISTRO_ID $DISTRO_VERSION${RESET}..."

    case "$DISTRO_ID" in
        ubuntu|debian|raspbian|linuxmint)
            info "Adding Docker repository and installing Docker for Debian/Ubuntu..."
            sudo apt update
            sudo apt install -y apt-transport-https ca-certificates curl gnupg lsb-release
            curl -fsSL https://download.docker.com/linux/"$DISTRO_ID"/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
            echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/$DISTRO_ID $DISTRO_CODENAME stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
            sudo apt update
            sudo apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
            sudo systemctl enable --now docker
            ;;

        fedora)
            info "Installing Docker from Fedora repositories..."
            sudo dnf remove -y podman-docker || true
            sudo dnf install -y dnf-plugins-core
            sudo dnf config-manager --add-repo https://download.docker.com/linux/fedora/docker-ce.repo
            sudo dnf install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
            sudo systemctl enable --now docker
            ;;

        centos|rhel|rocky|almalinux|ol|oracle)
            info "Adding Docker repository and installing Docker for CentOS/RHEL..."
            sudo yum install -y yum-utils
            sudo yum-config-manager --add-repo https://download.docker.com/linux/centos/docker-ce.repo
            sudo yum install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
            sudo systemctl enable --now docker
            ;;

        arch|manjaro|endeavouros|garuda|artix|arcolinux|antergos|chakra|kaos)
            info "Installing Docker with pacman (Arch-based)..."
            sudo pacman -Sy --noconfirm docker docker-compose
            sudo systemctl enable --now docker
            ;;

        opensuse*|suse|sles)
            info "Installing Docker with zypper..."
            sudo zypper install -y docker docker-compose
            sudo systemctl enable --now docker
            ;;

        alpine)
            info "Installing Docker with apk..."
            sudo apk add --update docker docker-compose
            sudo rc-update add docker boot
            sudo service docker start
            ;;

        *)
            error "Unsupported or unrecognized Linux distribution: $DISTRO_ID"
            exit 1
            ;;
    esac

    # Fix permissions for current user
    sudo chown -R "$USER":"$USER" "$HOME/.docker" || true
    sudo chmod -R g+rwx "$HOME/.docker" || true

    if ! groups "$USER" | grep -qw docker; then
        sudo usermod -aG docker "$USER"
        warn "User '$USER' has been added to the 'docker' group."
        error "IMPORTANT: Docker group permissions will NOT be active in this current session."
        error "Please LOG OUT completely and LOG BACK IN (or reboot the machine) then re-run this script."
        exit 1
    else
        success "User '$USER' is already in the 'docker' group."
    fi
}

install_docker_compose_standalone() {
    step "Installing Docker Compose standalone..."
    sudo curl -L "https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
    sudo chmod +x /usr/local/bin/docker-compose
    success "Docker Compose installed."
}

step "Checking for Docker..."
if ! command -v docker &>/dev/null; then
    warn "Docker not found. Starting installation..."
    install_docker
else
    success "Docker is already installed."
fi

step "Verifying Docker daemon access permissions for current user..."
if ! docker info &>/dev/null; then
    error "Permission denied: Cannot connect to Docker daemon."
    error "Your user ('$USER') does not have correct permissions for Docker in this session."
    info "Current groups for user: $(groups "$USER")"
    info "Docker socket permissions: $(ls -l /var/run/docker.sock 2>/dev/null || echo 'Socket not found')"
    error "Try running this command to refresh group membership in this terminal:"
    echo -e "${YELLOW}    newgrp docker${RESET}"
    error "Or LOG OUT completely and LOG BACK IN (or reboot the machine) then re-run this script."
    exit 1
else
    success "Docker daemon is accessible with current user permissions."
fi

step "Checking for Docker Compose..."
if ! docker compose version &>/dev/null; then
    if ! command -v docker-compose &>/dev/null; then
        warn "Docker Compose not found. Installing standalone version..."
        install_docker_compose_standalone
    else
        success "Docker Compose (standalone) is already installed."
    fi
else
    success "Docker Compose (plugin) is already installed."
fi

step "Ensuring Docker service is running..."
if [[ "$DISTRO_ID" == "alpine" ]]; then
    if ! sudo service docker status 2>/dev/null | grep -q 'status: started'; then
        info "Starting Docker service..."
        sudo service docker start || true
        sleep 5  # Give Docker time to start
    else
        success "Docker service is running."
    fi
else
    if ! sudo systemctl is-active --quiet docker; then
        info "Starting Docker service..."
        sudo systemctl start docker || true
        sleep 5  # Give Docker time to start
    else
        success "Docker service is running."
    fi
fi

# --- Portainer Installation/Update ---
UPDATE_PORTAINER=${1:-""}

step "Checking for existing Portainer container..."
if docker ps -a --format '{{.Names}}' | grep -q "^portainer$"; then
    if [[ "$UPDATE_PORTAINER" == "--update" ]]; then
        warn "Updating Portainer container..."
        docker stop portainer || true
        docker rm portainer || true
    else
        warn "Portainer container already exists. Use '--update' to force an update."
        exit 0
    fi
else
    info "No existing Portainer container found. Proceeding with installation."
fi

step "Pulling latest Portainer image..."
docker pull portainer/portainer-ce:latest

step "Starting Portainer container..."
if [[ "$DISTRO_ID" == "fedora" ]]; then
  docker run -d \
    --privileged \
    -p 8000:8000 -p 9443:9443 \
    --name portainer \
    --restart=always \
    -v /var/run/docker.sock:/var/run/docker.sock \
    -v portainer_data:/data \
    portainer/portainer-ce:latest
else
  docker run -d \
    -p 8000:8000 -p 9443:9443 \
    --name portainer \
    --restart=always \
    -v /var/run/docker.sock:/var/run/docker.sock \
    -v portainer_data:/data \
    portainer/portainer-ce:latest
fi

success "Portainer is now installed and running!"

# --- IP Detection ---
IP=$(hostname -I 2>/dev/null | awk '{print $1}' || ip route get 1 | awk '{print $7}' || echo "localhost")

echo -e "${BOLD}${CYAN}🌐 Access your Portainer dashboard at:${RESET} ${GREEN}https://$IP:9443${RESET}"
echo -e "${YELLOW}Note: You may need to accept a self-signed certificate in your browser for HTTPS.${RESET}"