#!/bin/bash

# ============================================================
# Docker & Docker Compose one-shot installer
# Target OS: Ubuntu 24.04.4 LTS
# Published by: helper.sh
# Script note: generated with Claude and curated by helper.sh
# ============================================================

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

DAEMON_JSON="/etc/docker/daemon.json"
DOCKER_REPO_FILE="/etc/apt/sources.list.d/docker.list"
DOCKER_GPG_KEY="/etc/apt/keyrings/docker.asc"
DOCKER_REPO_BASE="https://download.docker.com/linux/ubuntu"
DOCKER_REPO_MIRROR="https://mirrors.aliyun.com/docker-ce/linux/ubuntu"
REGISTRY_MIRROR="https://docker.1ms.run"
DAEMON_BACKUP=""
SELECTED_REPO_BASE=""

info()    { echo -e "${BLUE}[INFO]${NC}  $1"; }
success() { echo -e "${GREEN}[OK]${NC}    $1"; }
warning() { echo -e "${YELLOW}[WARN]${NC}  $1"; }
error()   { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

backup_daemon_json_once() {
  if [ -f "${DAEMON_JSON}" ] && [ -z "${DAEMON_BACKUP}" ]; then
    DAEMON_BACKUP="${DAEMON_JSON}.bak.$(date +%Y%m%d%H%M%S)"
    cp "${DAEMON_JSON}" "${DAEMON_BACKUP}"
    info "Backed up existing daemon.json to ${DAEMON_BACKUP}"
  fi
}

ensure_python3() {
  if ! command -v python3 >/dev/null 2>&1; then
    info "Installing python3 for safe daemon.json merge..."
    apt-get install -y -qq python3
  fi
}

merge_daemon_config() {
  local add_mirror="${1:-0}"

  ensure_python3
  backup_daemon_json_once

  python3 - "${DAEMON_JSON}" "${REGISTRY_MIRROR}" "${add_mirror}" <<'PY'
import json
import os
import sys

path, mirror, add_mirror = sys.argv[1], sys.argv[2], sys.argv[3] == "1"
data = {}

if os.path.exists(path):
    with open(path, "r", encoding="utf-8") as f:
        raw = f.read().strip()
    if raw:
        try:
            data = json.loads(raw)
        except json.JSONDecodeError as exc:
            print(f"[ERROR] daemon.json is not valid JSON: {exc}", file=sys.stderr)
            sys.exit(1)

if not isinstance(data, dict):
    print("[ERROR] daemon.json top-level structure must be a JSON object", file=sys.stderr)
    sys.exit(1)

data.setdefault("log-driver", "json-file")
log_opts = data.get("log-opts")
if not isinstance(log_opts, dict):
    log_opts = {}
data["log-opts"] = log_opts
log_opts.setdefault("max-size", "100m")
log_opts.setdefault("max-file", "3")

if add_mirror:
    mirrors = data.get("registry-mirrors")
    if not isinstance(mirrors, list):
        mirrors = []
    if mirror not in mirrors:
        mirrors.append(mirror)
    data["registry-mirrors"] = mirrors

with open(path, "w", encoding="utf-8") as f:
    json.dump(data, f, indent=2, ensure_ascii=False)
    f.write("\n")
PY
}

restart_docker() {
  systemctl daemon-reload
  systemctl restart docker
}

url_reachable() {
  local url="$1"
  curl -fsSL --connect-timeout 5 --max-time 15 -o /dev/null "${url}"
}

select_docker_repo_base() {
  local release_path="${DOCKER_REPO_BASE}/dists/${VERSION_CODENAME}/Release"
  local mirror_release_path="${DOCKER_REPO_MIRROR}/dists/${VERSION_CODENAME}/Release"

  if url_reachable "${release_path}"; then
    SELECTED_REPO_BASE="${DOCKER_REPO_BASE}"
    info "Using official Docker APT repository"
    return 0
  fi

  warning "Official Docker APT repository is unreachable, trying Aliyun mirror..."
  if url_reachable "${mirror_release_path}"; then
    SELECTED_REPO_BASE="${DOCKER_REPO_MIRROR}"
    success "Using Aliyun Docker CE mirror for package installation"
    return 0
  fi

  error "Neither the official Docker repository nor the Aliyun mirror is reachable"
}

configure_docker_repo() {
  local repo_base="$1"

  info "Adding Docker GPG key..."
  install -m 0755 -d /etc/apt/keyrings
  if ! curl -fsSL "${repo_base}/gpg" -o "${DOCKER_GPG_KEY}"; then
    if [ "${repo_base}" != "${DOCKER_REPO_MIRROR}" ]; then
      warning "Failed to download the Docker GPG key from the official source, switching to the Aliyun mirror..."
      SELECTED_REPO_BASE="${DOCKER_REPO_MIRROR}"
      configure_docker_repo "${DOCKER_REPO_MIRROR}"
      return 0
    fi
    error "Failed to download the Docker GPG key from ${repo_base}"
  fi
  chmod a+r "${DOCKER_GPG_KEY}"
  success "Docker GPG key installed"

  info "Adding Docker APT repository..."
  echo \
    "deb [arch=$(dpkg --print-architecture) signed-by=${DOCKER_GPG_KEY}] ${repo_base} $(. /etc/os-release && echo "${VERSION_CODENAME}") stable" \
    | tee "${DOCKER_REPO_FILE}" >/dev/null
  success "Docker repository configured"
}

check_global_registry_connectivity() {
  local code
  code="$(curl -sS -o /dev/null -w '%{http_code}' --connect-timeout 5 --max-time 10 https://registry-1.docker.io/v2/ || true)"
  case "${code}" in
    200|301|302|401|403)
      return 0
      ;;
    *)
      return 1
      ;;
  esac
}

ensure_registry_mirror_if_needed() {
  if check_global_registry_connectivity; then
    info "Docker Registry connectivity looks healthy. Keeping current pull settings."
    return 0
  fi

  warning "Docker Hub seems unreachable. Adding registry mirror ${REGISTRY_MIRROR}"
  merge_daemon_config 1
  restart_docker
  success "Registry mirror added and Docker restarted"
}

if [ "${EUID}" -ne 0 ]; then
  error "Run this script with root privileges: sudo bash $0"
fi

echo ""
echo -e "${BLUE}============================================================${NC}"
echo -e "${BLUE} Docker & Docker Compose one-shot installer                  ${NC}"
echo -e "${BLUE} Target OS: Ubuntu 24.04.4 LTS                              ${NC}"
echo -e "${BLUE} Published by: helper.sh                                    ${NC}"
echo -e "${BLUE}============================================================${NC}"
echo ""

info "Checking operating system..."
. /etc/os-release
if [[ "${ID}" != "ubuntu" ]]; then
  error "This script only supports Ubuntu. Current OS: ${ID}"
fi
if [[ "${VERSION_ID}" != "24.04" ]]; then
  warning "Current host is ${PRETTY_NAME}. This script is written for Ubuntu 24.04.4 LTS."
fi
success "Detected system: ${PRETTY_NAME}"

info "Removing older Docker packages if present..."
OLD_PKGS="docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc"
for pkg in ${OLD_PKGS}; do
  if dpkg -s "${pkg}" >/dev/null 2>&1; then
    apt-get remove -y "${pkg}" >/dev/null 2>&1 || true
    warning "Removed old package: ${pkg}"
  fi
done
success "Legacy package cleanup finished"

info "Updating package index and installing dependencies..."
apt-get update -qq
apt-get install -y -qq \
  ca-certificates \
  curl \
  gnupg \
  lsb-release \
  python3
success "Dependencies installed"

select_docker_repo_base
configure_docker_repo "${SELECTED_REPO_BASE}"

info "Refreshing package index..."
if ! apt-get update -qq; then
  if [ "${SELECTED_REPO_BASE}" != "${DOCKER_REPO_MIRROR}" ]; then
    warning "apt update failed against the official Docker repository, switching to the Aliyun mirror..."
    SELECTED_REPO_BASE="${DOCKER_REPO_MIRROR}"
    configure_docker_repo "${SELECTED_REPO_BASE}"
    apt-get update -qq
  else
    error "apt update failed after switching to the Docker mirror"
  fi
fi

info "Installing Docker Engine and Docker Compose..."
apt-get install -y \
  docker-ce \
  docker-ce-cli \
  containerd.io \
  docker-buildx-plugin \
  docker-compose-plugin
success "Docker packages installed"

info "Starting and enabling Docker..."
systemctl start docker
systemctl enable docker
success "Docker service is running and enabled"

REAL_USER="${SUDO_USER:-$USER}"
if [ -n "${REAL_USER}" ] && [ "${REAL_USER}" != "root" ]; then
  info "Adding '${REAL_USER}' to the docker group..."
  usermod -aG docker "${REAL_USER}"
  success "User '${REAL_USER}' added to docker group (takes effect after re-login)"
else
  warning "Current user is root. Skipping docker group assignment."
fi

info "Merging baseline Docker daemon settings..."
mkdir -p /etc/docker
merge_daemon_config 0
restart_docker
success "Baseline daemon configuration applied"

info "Checking Docker Hub reachability..."
ensure_registry_mirror_if_needed

info "Verifying Docker installation..."
DOCKER_VERSION="$(docker --version)"
success "Docker version: ${DOCKER_VERSION}"

info "Verifying Docker Compose installation..."
COMPOSE_VERSION="$(docker compose version)"
success "Docker Compose version: ${COMPOSE_VERSION}"

info "Running hello-world test container..."
if docker run --rm hello-world >/dev/null 2>&1; then
  success "hello-world container ran successfully"
else
  warning "hello-world failed. Applying registry mirror and retrying..."
  merge_daemon_config 1
  restart_docker
  if docker run --rm hello-world >/dev/null 2>&1; then
    success "hello-world succeeded after enabling the registry mirror"
  else
    warning "hello-world still failed. Check network, DNS, or Docker service health."
  fi
fi

echo ""
echo -e "${GREEN}============================================================${NC}"
echo -e "${GREEN} Docker installation completed                               ${NC}"
echo -e "${GREEN}============================================================${NC}"
echo ""
echo -e "  ${BLUE}Docker version:${NC}         ${DOCKER_VERSION}"
echo -e "  ${BLUE}Docker Compose version:${NC} ${COMPOSE_VERSION}"
echo ""
echo -e "  ${YELLOW}Useful commands:${NC}"
echo -e "    docker ps                    # Show running containers"
echo -e "    docker images                # Show local images"
echo -e "    docker compose up -d         # Start a Compose stack in detached mode"
echo -e "    systemctl status docker      # Check Docker service status"
echo ""
if grep -q "${REGISTRY_MIRROR}" "${DAEMON_JSON}" 2>/dev/null; then
  echo -e "  ${YELLOW}Registry mirror:${NC} ${REGISTRY_MIRROR}"
fi
if [ -n "${DAEMON_BACKUP}" ]; then
  echo -e "  ${YELLOW}daemon.json backup:${NC} ${DAEMON_BACKUP}"
fi
if [ -n "${REAL_USER}" ] && [ "${REAL_USER}" != "root" ]; then
  echo -e "  ${YELLOW}Note:${NC} '${REAL_USER}' is now in the docker group."
  echo -e "        Re-login or run ${BLUE}newgrp docker${NC} to use Docker without sudo."
  echo ""
fi