#!/usr/bin/env bash
#
# install-docker.sh — Detect or install Docker + Docker Compose plugin
# Production-ready. Idempotent. Verifies every step.
#
set -o pipefail

# ---------- Colors ----------
if [[ -t 1 ]]; then
    C_RESET=$'\e[0m'
    C_ORANGE=$'\e[38;5;208m'
    C_GREEN=$'\e[32m'
    C_RED=$'\e[31m'
    C_CYAN=$'\e[36m'
    C_BOLD=$'\e[1m'
else
    C_RESET='' C_ORANGE='' C_GREEN='' C_RED='' C_CYAN='' C_BOLD=''
fi

proc()  { printf "%s[PROC]%s %s\n" "$C_ORANGE" "$C_RESET" "$*"; }
ok()    { printf "%s[ OK ]%s %s\n" "$C_GREEN"  "$C_RESET" "$*"; }
fail()  { printf "%s[FAIL]%s %s\n" "$C_RED"    "$C_RESET" "$*" >&2; }
info()  { printf "%s[INFO]%s %s\n" "$C_CYAN"   "$C_RESET" "$*"; }

die() { fail "$*"; exit 1; }

# ---------- Banner ----------
banner() {
    clear
    printf "%s%s" "$C_CYAN" "$C_BOLD"
    cat <<'EOF'
 ___           ___                      _
|_ _|_ __  ___/ _ \ _ __ ___  _ __  (_) __ _
 | || '_ \/ __| | | | '_ ` _ \| '_ \ | |/ _` |
 | || | | \__ \ |_| | | | | | | | | || | (_| |
|___|_| |_|___/\___/|_| |_| |_|_| |_||_|\__,_|

EOF
    cat <<EOF
╔══════════════════════════════════════════════════════════════╗
║                                                              ║
║           DOCKER  INSTALLER  &  VERIFIER  v1.1               ║
║              Production-ready • Multi-distro                 ║
║                                                              ║
╚══════════════════════════════════════════════════════════════╝
${C_RESET}
EOF
}

# ---------- Y/N prompt ----------
ask_yn() {
    local prompt=$1 default=${2:-N} reply
    local hint="[y/N]"; [[ "$default" == "Y" ]] && hint="[Y/n]"
    while true; do
        printf "%s[ ?? ]%s %s %s " "$C_ORANGE" "$C_RESET" "$prompt" "$hint"
        read -r reply </dev/tty || reply=""
        [[ -z "$reply" ]] && reply="$default"
        case "$reply" in
            Y*|y*) return 0 ;;
            N*|n*) return 1 ;;
        esac
    done
}

# ---------- Sudo wrapper ----------
SUDO=""
setup_sudo() {
    if [[ $EUID -eq 0 ]]; then
        SUDO=""
        ok "Running as root"
    elif command -v sudo >/dev/null 2>&1; then
        SUDO="sudo"
        ok "sudo available — will use for privileged operations"
    else
        die "Not root and sudo not installed. Install sudo or run as root."
    fi
}

# ---------- Dependency check ----------
ensure_dep() {
    local cmd=$1 pkg=${2:-$1}
    if command -v "$cmd" >/dev/null 2>&1; then
        ok "Dependency present: $cmd"
        return 0
    fi
    proc "Installing missing dependency: $pkg"
    case "$PKG_MGR" in
        apt)    $SUDO apt-get update -qq && $SUDO apt-get install -y -qq "$pkg" ;;
        dnf)    $SUDO dnf install -y -q "$pkg" ;;
        yum)    $SUDO yum install -y -q "$pkg" ;;
        pacman) $SUDO pacman -Sy --noconfirm --needed "$pkg" ;;
        zypper) $SUDO zypper --non-interactive install "$pkg" ;;
        apk)    $SUDO apk add --quiet "$pkg" ;;
        *)      die "Unknown package manager; cannot install $pkg" ;;
    esac
    command -v "$cmd" >/dev/null 2>&1 || die "Failed to install $pkg"
    ok "Installed: $pkg"
}

# ---------- OS detection ----------
detect_os() {
    proc "Detecting operating system..."
    [[ "$(uname -s)" == "Linux" ]] || die "This script supports Linux only. For macOS/Windows use Docker Desktop."

    [[ -r /etc/os-release ]] || die "/etc/os-release not found — unsupported distribution."
    # shellcheck disable=SC1091
    . /etc/os-release
    OS_ID="${ID:-unknown}"
    OS_LIKE="${ID_LIKE:-}"
    OS_VER="${VERSION_ID:-}"
    OS_CODENAME="${VERSION_CODENAME:-}"
    OS_NAME="${PRETTY_NAME:-$OS_ID}"

    case "$OS_ID $OS_LIKE" in
        *debian*|*ubuntu*) PKG_MGR="apt"   ; OS_FAMILY="debian" ;;
        *rhel*|*fedora*|*centos*|*rocky*|*almalinux*)
            if   command -v dnf >/dev/null 2>&1; then PKG_MGR="dnf"
            else PKG_MGR="yum"; fi
            OS_FAMILY="rhel" ;;
        *arch*)            PKG_MGR="pacman"; OS_FAMILY="arch"   ;;
        *suse*|*opensuse*) PKG_MGR="zypper"; OS_FAMILY="suse"   ;;
        *alpine*)          PKG_MGR="apk"   ; OS_FAMILY="alpine" ;;
        *) die "Unsupported distribution: $OS_ID" ;;
    esac

    ok "OS: $OS_NAME (family: $OS_FAMILY, pkg: $PKG_MGR)"
}

# ---------- Docker check ----------
check_docker() {
    proc "Checking if Docker is installed..."
    if ! command -v docker >/dev/null 2>&1; then
        info "Docker is NOT installed."
        return 1
    fi

    local ver
    ver=$(docker --version 2>/dev/null) || { fail "docker binary present but not working"; return 1; }
    ok "Docker binary found: $ver"

    proc "Verifying Docker daemon..."
    if $SUDO docker info >/dev/null 2>&1; then
        ok "Docker daemon is running and reachable"
    else
        fail "Docker installed but daemon not reachable"
        return 1
    fi

    proc "Checking Docker Compose plugin..."
    if $SUDO docker compose version >/dev/null 2>&1; then
        ok "Docker Compose: $($SUDO docker compose version 2>/dev/null)"
    else
        info "Docker Compose plugin missing — will install"
        return 2
    fi

    return 0
}

# ---------- Check for updates to Docker packages ----------
# Populates global array UPDATE_PKGS with names of Docker-related packages that have pending updates.
UPDATE_PKGS=()
check_docker_updates() {
    UPDATE_PKGS=()
    proc "Checking for available updates to Docker components..."

    local -a watch=(docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin \
                    docker docker-compose docker-cli-compose docker-buildx)

    case "$PKG_MGR" in
        apt)
            $SUDO apt-get update -qq >/dev/null 2>&1 || { fail "apt-get update failed"; return 1; }
            local line pkg
            while IFS= read -r line; do
                pkg="${line%%/*}"
                for w in "${watch[@]}"; do
                    [[ "$pkg" == "$w" ]] && UPDATE_PKGS+=("$pkg") && break
                done
            done < <(apt list --upgradable 2>/dev/null | tail -n +2)
            ;;
        dnf|yum)
            local out
            out=$($SUDO $PKG_MGR -q check-update 2>/dev/null || true)
            for w in "${watch[@]}"; do
                grep -q "^${w}\." <<<"$out" && UPDATE_PKGS+=("$w")
            done
            ;;
        pacman)
            $SUDO pacman -Sy >/dev/null 2>&1 || true
            local out
            out=$(pacman -Qu 2>/dev/null || true)
            for w in "${watch[@]}"; do
                grep -q "^${w} " <<<"$out" && UPDATE_PKGS+=("$w")
            done
            ;;
        zypper)
            local out
            out=$($SUDO zypper -q list-updates 2>/dev/null || true)
            for w in "${watch[@]}"; do
                grep -q " ${w} " <<<"$out" && UPDATE_PKGS+=("$w")
            done
            ;;
        apk)
            $SUDO apk update -q >/dev/null 2>&1 || true
            local out
            out=$(apk version -l '<' 2>/dev/null | tail -n +2 || true)
            for w in "${watch[@]}"; do
                grep -q "^${w}-" <<<"$out" && UPDATE_PKGS+=("$w")
            done
            ;;
    esac

    if ((${#UPDATE_PKGS[@]} == 0)); then
        ok "All Docker components are up to date."
        return 0
    fi

    info "Updates available for the following Docker components:"
    local p
    for p in "${UPDATE_PKGS[@]}"; do
        printf "       %s•%s %s\n" "$C_ORANGE" "$C_RESET" "$p"
    done
    return 0
}

# ---------- Apply updates to the listed Docker packages ----------
apply_docker_updates() {
    ((${#UPDATE_PKGS[@]} > 0)) || return 0
    proc "Updating Docker components: ${UPDATE_PKGS[*]}"
    case "$PKG_MGR" in
        apt)    $SUDO apt-get install -y -qq --only-upgrade "${UPDATE_PKGS[@]}" ;;
        dnf)    $SUDO dnf upgrade -y -q "${UPDATE_PKGS[@]}" ;;
        yum)    $SUDO yum update  -y -q "${UPDATE_PKGS[@]}" ;;
        pacman) $SUDO pacman -S --noconfirm --needed "${UPDATE_PKGS[@]}" ;;
        zypper) $SUDO zypper --non-interactive update "${UPDATE_PKGS[@]}" ;;
        apk)    $SUDO apk upgrade "${UPDATE_PKGS[@]}" ;;
    esac || { fail "Update failed"; return 1; }

    proc "Verifying Docker still works after update..."
    $SUDO docker info >/dev/null 2>&1 || die "Docker daemon broken after update"
    $SUDO docker compose version >/dev/null 2>&1 || die "Docker Compose broken after update"
    ok "Docker updated successfully: $(docker --version)"
    ok "Compose: $($SUDO docker compose version)"
}

# ---------- OS security updates (NEVER distro upgrade) ----------
os_update() {
    proc "Fetching available OS package updates (security + regular — NO distro upgrade)..."
    local rc=0
    case "$PKG_MGR" in
        apt)
            export DEBIAN_FRONTEND=noninteractive
            if [[ -n "$SUDO" ]]; then
                sudo -E apt-get update -qq || { fail "apt-get update failed"; return 1; }
                sudo -E apt-get upgrade -y -qq \
                    -o Dpkg::Options::="--force-confdef" \
                    -o Dpkg::Options::="--force-confold"
            else
                apt-get update -qq || { fail "apt-get update failed"; return 1; }
                apt-get upgrade -y -qq \
                    -o Dpkg::Options::="--force-confdef" \
                    -o Dpkg::Options::="--force-confold"
            fi
            rc=$?
            ;;
        dnf)    $SUDO dnf  upgrade --refresh -y -q ; rc=$? ;;
        yum)    $SUDO yum  update  -y -q           ; rc=$? ;;
        pacman) $SUDO pacman -Syu --noconfirm      ; rc=$? ;;
        zypper) $SUDO zypper --non-interactive refresh && $SUDO zypper --non-interactive update ; rc=$? ;;
        apk)    $SUDO apk update -q && $SUDO apk upgrade ; rc=$? ;;
    esac
    if ((rc == 0)); then
        ok "OS packages updated (distribution version was NOT changed)."
        return 0
    else
        fail "OS update encountered errors — review output above."
        return 1
    fi
}

# ---------- Install Docker ----------
install_docker_debian() {
    proc "Installing Docker on Debian/Ubuntu..."
    ensure_dep curl
    ensure_dep gpg gnupg
    ensure_dep lsb_release lsb-release || true

    $SUDO install -m 0755 -d /etc/apt/keyrings
    local repo_id="$OS_ID"
    [[ "$OS_ID" != "ubuntu" && "$OS_ID" != "debian" ]] && {
        case "$OS_LIKE" in *ubuntu*) repo_id="ubuntu";; *) repo_id="debian";; esac
    }

    curl -fsSL "https://download.docker.com/linux/${repo_id}/gpg" \
        | $SUDO gpg --dearmor --yes -o /etc/apt/keyrings/docker.gpg
    $SUDO chmod a+r /etc/apt/keyrings/docker.gpg

    local codename="${OS_CODENAME:-$(lsb_release -cs 2>/dev/null || echo stable)}"
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/${repo_id} ${codename} stable" \
        | $SUDO tee /etc/apt/sources.list.d/docker.list >/dev/null

    $SUDO apt-get update -qq
    $SUDO apt-get install -y -qq docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
}

install_docker_rhel() {
    proc "Installing Docker on RHEL/Fedora/Rocky/Alma..."
    ensure_dep curl
    $SUDO $PKG_MGR install -y -q "${PKG_MGR}-plugins-core" 2>/dev/null || true
    local repo_id="centos"
    [[ "$OS_ID" == "fedora" ]] && repo_id="fedora"
    $SUDO $PKG_MGR config-manager --add-repo "https://download.docker.com/linux/${repo_id}/docker-ce.repo" \
        || $SUDO curl -fsSL "https://download.docker.com/linux/${repo_id}/docker-ce.repo" -o /etc/yum.repos.d/docker-ce.repo
    $SUDO $PKG_MGR install -y -q docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
}

install_docker_arch() {
    proc "Installing Docker on Arch Linux..."
    $SUDO pacman -Sy --noconfirm --needed docker docker-compose docker-buildx
}

install_docker_suse() {
    proc "Installing Docker on openSUSE/SLES..."
    $SUDO zypper --non-interactive install docker docker-compose
}

install_docker_alpine() {
    proc "Installing Docker on Alpine..."
    $SUDO apk add --quiet docker docker-cli-compose
    $SUDO rc-update add docker boot 2>/dev/null || true
}

install_docker() {
    case "$OS_FAMILY" in
        debian) install_docker_debian ;;
        rhel)   install_docker_rhel   ;;
        arch)   install_docker_arch   ;;
        suse)   install_docker_suse   ;;
        alpine) install_docker_alpine ;;
        *) die "No installer for OS family: $OS_FAMILY" ;;
    esac

    proc "Enabling and starting Docker service..."
    if command -v systemctl >/dev/null 2>&1; then
        $SUDO systemctl enable --now docker \
            || die "Failed to enable/start docker.service"
    elif command -v rc-service >/dev/null 2>&1; then
        $SUDO rc-service docker start || die "Failed to start docker"
    fi

    proc "Verifying installation..."
    command -v docker >/dev/null 2>&1 || die "docker binary still missing after install"
    $SUDO docker info >/dev/null 2>&1 || die "docker daemon not reachable after install"
    $SUDO docker compose version >/dev/null 2>&1 || die "docker compose plugin not available after install"

    ok "Docker:         $(docker --version)"
    ok "Docker Compose: $($SUDO docker compose version)"

    proc "Running hello-world container as final sanity check..."
    if $SUDO docker run --rm hello-world >/dev/null 2>&1; then
        ok "hello-world container ran successfully"
    else
        fail "hello-world failed — check network/daemon, but binaries are installed"
    fi

    if [[ $EUID -ne 0 ]]; then
        proc "Adding user '$(whoami)' to docker group..."
        $SUDO usermod -aG docker "$(whoami)" && \
            ok "User added to docker group — log out and back in to use docker without sudo"
    fi
}

# ---------- Main ----------
main() {
    banner
    setup_sudo
    detect_os

    set +e
    check_docker
    local rc=$?
    set -e

    case $rc in
        0)
            ok "Docker and Docker Compose are already installed and working."
            # Check for available updates
            check_docker_updates || true
            if ((${#UPDATE_PKGS[@]} > 0)); then
                if ask_yn "Download and install these Docker updates now?" Y; then
                    apply_docker_updates
                else
                    info "Skipping Docker updates per user choice."
                fi
            fi
            ;;
        2)
            info "Docker is installed but the Compose plugin is missing."
            if ask_yn "Install the missing Docker Compose plugin now?" Y; then
                case "$OS_FAMILY" in
                    debian) $SUDO apt-get update -qq && $SUDO apt-get install -y -qq docker-compose-plugin ;;
                    rhel)   $SUDO $PKG_MGR install -y -q docker-compose-plugin ;;
                    arch)   $SUDO pacman -Sy --noconfirm --needed docker-compose ;;
                    suse)   $SUDO zypper --non-interactive install docker-compose ;;
                    alpine) $SUDO apk add --quiet docker-cli-compose ;;
                esac
                $SUDO docker compose version >/dev/null 2>&1 \
                    && ok "Compose plugin installed: $($SUDO docker compose version)" \
                    || die "Compose plugin install failed"
            else
                info "Skipping Compose plugin install per user choice."
            fi
            # Also offer Docker component updates
            check_docker_updates || true
            if ((${#UPDATE_PKGS[@]} > 0)) && ask_yn "Download and install these Docker updates now?" Y; then
                apply_docker_updates
            fi
            ;;
        1)
            info "Docker is NOT installed."
            if ask_yn "Install Docker CE + Compose plugin from the official repository now?" Y; then
                install_docker
                ok "Installation complete."
            else
                info "User declined Docker installation. Exiting."
                exit 0
            fi
            ;;
    esac

    # Final OS update prompt (never distro upgrade)
    echo
    info "Final step: operating system package updates."
    info "This will apply security + regular updates but will NEVER upgrade the distribution version."
    if ask_yn "Run OS package updates now?" N; then
        os_update
    else
        info "Skipping OS updates per user choice."
    fi

    echo
    ok "All done. Enjoy your Docker setup!"
    exit 0
}

main "$@"