#!/bin/bash
# =============================================================================
# 🐳 StianNOR — PORTAINER INSTALLER v2.1
# =============================================================================
set -euo pipefail

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'
LINE="────────────────────────────────────────────────────────────────────────────"

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

FORCE_UPDATE=false
[[ "${1:-}" == "--update" ]] && FORCE_UPDATE=true

# =============================================================================
# SUDO — cache + keepalive
# =============================================================================
step "Checking sudo access"
sudo -v || die "sudo required"
( while true; do sudo -n true; sleep 55; done ) &
KEEPALIVE_PID=$!
trap 'kill "$KEEPALIVE_PID" 2>/dev/null || true' EXIT INT TERM

# =============================================================================
# DETECT DISTRO
# =============================================================================
step "Detecting Linux distribution"
[[ ! -f /etc/os-release ]] && die "Cannot detect distro — /etc/os-release missing"
source /etc/os-release
DISTRO_ID="${ID,,}"
DISTRO_CODENAME="${VERSION_CODENAME:-}"
[[ -z "$DISTRO_CODENAME" ]] && DISTRO_CODENAME=$(lsb_release -cs 2>/dev/null || echo "")
[[ -z "$DISTRO_CODENAME" && "$DISTRO_ID" == "ubuntu" ]] && DISTRO_CODENAME="jammy"
[[ -z "$DISTRO_CODENAME" && "$DISTRO_ID" == "debian" ]] && DISTRO_CODENAME="bookworm"

# Map to distro family
DISTRO_FAMILY="$DISTRO_ID"
case "$DISTRO_ID" in
    ubuntu|debian|linuxmint|raspbian|pop) DISTRO_FAMILY="debian" ;;
    fedora)                               DISTRO_FAMILY="fedora" ;;
    centos|rhel|rocky|almalinux|ol)       DISTRO_FAMILY="rhel"   ;;
    arch|manjaro|endeavouros|garuda|artix|arcolinux) DISTRO_FAMILY="arch" ;;
    opensuse*|suse|sles)                  DISTRO_FAMILY="suse"   ;;
    alpine)                               DISTRO_FAMILY="alpine" ;;
    *)
        # Try ID_LIKE fallback
        case "${ID_LIKE:-}" in
            *debian*|*ubuntu*) DISTRO_FAMILY="debian" ;;
            *rhel*|*fedora*)   DISTRO_FAMILY="fedora" ;;
            *arch*)            DISTRO_FAMILY="arch"   ;;
            *suse*)            DISTRO_FAMILY="suse"   ;;
        esac ;;
esac
success "Detected: ${NAME} (family: ${DISTRO_FAMILY})"

# =============================================================================
# INSTALL DOCKER
# =============================================================================
install_docker() {
    step "Installing Docker CE"
    case "$DISTRO_FAMILY" in

        debian)
            sudo apt-get update -qq
            sudo apt-get install -y -qq ca-certificates curl gnupg lsb-release
            sudo install -m0755 -d /etc/apt/keyrings
            curl -fsSL "https://download.docker.com/linux/${DISTRO_ID}/gpg" \
                | sudo gpg --dearmor --yes -o /etc/apt/keyrings/docker.gpg
            sudo chmod a+r /etc/apt/keyrings/docker.gpg
            echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
https://download.docker.com/linux/${DISTRO_ID} ${DISTRO_CODENAME} stable" \
                | sudo tee /etc/apt/sources.list.d/docker.list >/dev/null
            sudo apt-get update -qq
            sudo apt-get install -y \
                docker-ce docker-ce-cli containerd.io \
                docker-buildx-plugin docker-compose-plugin
            ;;

        fedora)
            sudo dnf -y remove podman-docker 2>/dev/null || true
            sudo dnf -y install dnf-plugins-core
            sudo dnf config-manager \
                --add-repo https://download.docker.com/linux/fedora/docker-ce.repo
            sudo dnf -y install \
                docker-ce docker-ce-cli containerd.io \
                docker-buildx-plugin docker-compose-plugin
            ;;

        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
            ;;

        arch)
            sudo pacman -Sy --noconfirm --needed docker docker-compose
            ;;

        suse)
            sudo zypper install -y docker docker-compose
            ;;

        alpine)
            sudo apk add --update docker docker-compose
            sudo rc-update add docker boot
            sudo service docker start
            success "Docker installed (Alpine)"
            return
            ;;

        *)
            die "Unsupported distro: $DISTRO_ID — see https://docs.docker.com/engine/install/"
            ;;
    esac

    sudo systemctl enable --now docker
    success "Docker installed and started"
}

# =============================================================================
# FIX DOCKER SOCKET PERMISSIONS
# The real fix for "permission denied on /var/run/docker.sock":
#   1. Add user to docker group
#   2. Fix socket group ownership so the group can actually access it
#   3. Tell the user to run `newgrp docker` in their shell — we cannot
#      push group changes to the parent shell process, that's a kernel limit.
# =============================================================================
fix_docker_permissions() {
    step "Fixing Docker permissions for $USER"

    # Ensure docker group exists
    getent group docker >/dev/null 2>&1 || sudo groupadd docker

    # Add user to docker group
    if ! groups "$USER" | grep -qw docker; then
        sudo usermod -aG docker "$USER"
        success "Added $USER to docker group"
    else
        success "$USER already in docker group"
    fi

    # Fix socket group ownership and permissions
    # This is what actually makes docker work without sudo RIGHT NOW
    if [[ -S /var/run/docker.sock ]]; then
        sudo chown root:docker /var/run/docker.sock
        sudo chmod 660 /var/run/docker.sock
        success "Fixed /var/run/docker.sock permissions (root:docker 660)"
    fi

    # Fix ~/.docker ownership
    sudo mkdir -p "$HOME/.docker"
    sudo chown -R "$USER:$USER" "$HOME/.docker"
    success "Fixed ~/.docker ownership"

    # Make socket permissions persist across reboots via systemd override
    OVERRIDE_DIR="/etc/systemd/system/docker.socket.d"
    if [[ -d /etc/systemd/system ]] && ! [[ -f "$OVERRIDE_DIR/permissions.conf" ]]; then
        sudo mkdir -p "$OVERRIDE_DIR"
        sudo tee "$OVERRIDE_DIR/permissions.conf" >/dev/null <<'OVERRIDE'
[Socket]
SocketMode=0660
SocketGroup=docker
OVERRIDE
        sudo systemctl daemon-reload
        sudo systemctl restart docker.socket docker 2>/dev/null || true
        success "Persistent socket permission override installed"
    fi
}

# =============================================================================
# ACTIVATE GROUP IN CURRENT SESSION
# sg/newgrp cannot push groups back to the parent shell.
# Best we can do: detect if THIS script process already has docker group,
# and if not, use `sg` to run the rest of the script with the group active.
# =============================================================================
activate_docker_group() {
    # Check if docker group is already active in this process
    if id -nG "$USER" 2>/dev/null | grep -qw docker && \
       [[ "$(stat -c '%G' /var/run/docker.sock 2>/dev/null)" == "docker" ]]; then
        # Try running docker without sudo
        if docker info >/dev/null 2>&1; then
            success "Docker accessible without sudo ✓"
            DOCKER_CMD="docker"
            return 0
        fi
    fi

    # Socket fix applied above — try again with newgrp trick
    # We use `sg docker` to get the group active for the REST of this script
    if [[ "${_DOCKER_GROUP_EXEC:-}" != "1" ]]; then
        export _DOCKER_GROUP_EXEC=1
        info "Activating docker group for this session via sg..."
        exec sg docker -c "export _DOCKER_GROUP_EXEC=1; bash '$0' ${1:-}"
    fi

    # If still failing, use sudo docker for the rest of the script
    if sudo docker info >/dev/null 2>&1; then
        warn "Using sudo docker for this session."
        warn "To use docker without sudo in YOUR terminal, run:"
        warn "  newgrp docker"
        warn "Or log out and back in."
        DOCKER_CMD="sudo docker"
    else
        die "Docker daemon is not accessible even with sudo. Check: sudo systemctl status docker"
    fi
}

# =============================================================================
# MAIN — INSTALL IF NEEDED
# =============================================================================
if ! command -v docker >/dev/null 2>&1; then
    install_docker
else
    success "Docker already installed: $(docker --version)"
fi

fix_docker_permissions
activate_docker_group

# =============================================================================
# DOCKER COMPOSE CHECK
# =============================================================================
step "Checking Docker Compose"
if $DOCKER_CMD compose version >/dev/null 2>&1; then
    success "Docker Compose plugin: $($DOCKER_CMD compose version --short 2>/dev/null || echo installed)"
elif command -v docker-compose >/dev/null 2>&1; then
    success "Docker Compose standalone: $(docker-compose --version)"
else
    info "Installing Docker Compose standalone..."
    sudo curl -fsSL \
        "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: $(/usr/local/bin/docker-compose --version)"
fi

# =============================================================================
# ENSURE DOCKER IS RUNNING
# =============================================================================
step "Ensuring Docker service is running"
if [[ "$DISTRO_FAMILY" == "alpine" ]]; then
    sudo service docker status 2>/dev/null | grep -q started \
        || { sudo service docker start; sleep 3; }
else
    if ! sudo systemctl is-active --quiet docker; then
        sudo systemctl start docker
        sleep 3
    fi
fi
success "Docker service running"

# =============================================================================
# PORTAINER
# =============================================================================
step "Installing Portainer CE"

PORTAINER_EXISTS=false
$DOCKER_CMD ps -a --format '{{.Names}}' 2>/dev/null | grep -q "^portainer$" \
    && PORTAINER_EXISTS=true

if [[ "$PORTAINER_EXISTS" == true ]]; then
    if [[ "$FORCE_UPDATE" == true ]]; then
        info "Removing existing Portainer for update..."
        $DOCKER_CMD stop portainer 2>/dev/null || true
        $DOCKER_CMD rm portainer 2>/dev/null || true
    else
        PORTAINER_STATUS=$($DOCKER_CMD inspect --format '{{.State.Status}}' portainer 2>/dev/null || echo "unknown")
        success "Portainer already installed (status: $PORTAINER_STATUS)"
        info "Run with --update to force reinstall"
        IP=$(hostname -I 2>/dev/null | awk '{print $1}')
        [[ -z "$IP" ]] && IP=$(ip route get 1.1.1.1 2>/dev/null | awk '{print $7; exit}')
        [[ -z "$IP" ]] && IP="localhost"
        echo -e "\n${LINE}"
        echo -e "  ${BOLD}${CYAN}🌐 Portainer:${RESET} ${GREEN}https://${IP}:9443${RESET}"
        echo -e "${LINE}"
        exit 0
    fi
fi

info "Pulling latest Portainer CE image..."
$DOCKER_CMD pull portainer/portainer-ce:latest

# SELinux detection
SELINUX_FLAG=""
if command -v getenforce >/dev/null 2>&1 && [[ "$(getenforce 2>/dev/null)" == "Enforcing" ]]; then
    SELINUX_FLAG="--security-opt label=disable"
    info "SELinux enforcing — adding label=disable"
fi

info "Starting Portainer container..."
$DOCKER_CMD run -d \
    $SELINUX_FLAG \
    -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

success "Portainer container started"

# =============================================================================
# DONE
# =============================================================================
IP=$(hostname -I 2>/dev/null | awk '{print $1}')
[[ -z "$IP" ]] && IP=$(ip route get 1.1.1.1 2>/dev/null | awk '{print $7; exit}')
[[ -z "$IP" ]] && IP="localhost"

# Check if docker still needs newgrp in the user's shell
NEEDS_NEWGRP=false
if ! docker info >/dev/null 2>&1; then
    NEEDS_NEWGRP=true
fi

echo ""
echo -e "${LINE}"
echo -e "${BOLD}${GREEN}  ✅ Portainer CE installed successfully!${RESET}"
echo -e "${LINE}"
echo -e "  ${CYAN}Dashboard:${RESET}  ${GREEN}https://${IP}:9443${RESET}"
echo -e "  ${CYAN}Tunnel:${RESET}     ${BLUE}http://${IP}:8000${RESET}"
echo -e "  ${YELLOW}Note:${RESET} Accept the self-signed certificate on first visit."
echo -e "  ${YELLOW}Note:${RESET} Create your admin account within 5 minutes of starting."
echo -e "${LINE}"

if [[ "$NEEDS_NEWGRP" == true ]]; then
    echo ""
    echo -e "${LINE}"
    echo -e "  ${BOLD}${YELLOW}⚠  ONE MORE STEP — activate docker group in your shell:${RESET}"
    echo -e "${LINE}"
    echo -e "  Docker is installed and working, but your CURRENT terminal"
    echo -e "  session doesn't have the docker group active yet."
    echo -e ""
    echo -e "  Run this in your terminal now:"
    echo -e "  ${BOLD}${CYAN}  newgrp docker${RESET}"
    echo -e ""
    echo -e "  Or simply log out and back in."
    echo -e "  After that, ${CYAN}docker ps${RESET} will work without sudo."
    echo -e "${LINE}"
fi
echo ""