#!/usr/bin/env bash
# ==============================================================================
# Metadata
# ==============================================================================
# Author      : Smith Wang (Refactor by ChatGPT)
# Version     : 2.0.0
# License     : MIT
# Description : Configure Docker APT repository (mirror) and install Docker on
#               Ubuntu (18.04/20.04/22.04/24.04) with robust offline handling.
#
# Modules     :
#   - Logging & Error Handling
#   - Environment & Dependency Checks
#   - Public Network Reachability Detection
#   - Docker GPG Key Installation (Online/Offline)
#   - Docker APT Repo Configuration
#   - Docker Installation & Service Setup
#
# Notes       :
#   - This script DOES NOT modify Ubuntu APT sources (/etc/apt/sources.list)
#   - This script DOES NOT set APT proxy (assumed handled elsewhere)
#   - If public network is NOT reachable and local GPG key is missing, script
#     will NOT proceed (per your requirement).
#
# ShellCheck  : Intended clean for bash v5+ with: shellcheck -x <script>
# ==============================================================================

set -euo pipefail

# ==============================================================================
# Global Constants
# ==============================================================================
readonly SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_VERSION="2.0.0"

# Default mirror for Docker repo (you asked: only focus on docker source)
readonly DEFAULT_DOCKER_APT_MIRROR="https://mirrors.aliyun.com/docker-ce/linux/ubuntu"

# Default keyring location (recommended by modern Ubuntu)
readonly DEFAULT_KEYRING_PATH="/etc/apt/keyrings/docker.gpg"

# Exit codes
readonly EC_OK=0
readonly EC_GENERAL=1
readonly EC_UNSUPPORTED_OS=10
readonly EC_DEPENDENCY=11
readonly EC_OFFLINE_NO_KEY=20
readonly EC_APT_FAILURE=30

# ==============================================================================
# Configurable Variables (Environment Overrides)
# ==============================================================================
# You may export these before running:
#   DOCKER_VERSION="20.10"         # or "20.10.15" (optional)
#   DOCKER_APT_MIRROR="https://..."
#   DOCKER_KEYRING_PATH="/root/wdd/docker.gpg"
#   LOCAL_DOCKER_GPG="/path/to/docker.gpg" (optional)
#   LOG_LEVEL="DEBUG|INFO|WARN|ERROR"
DOCKER_VERSION="${DOCKER_VERSION:-20.10}"
DOCKER_APT_MIRROR="${DOCKER_APT_MIRROR:-$DEFAULT_DOCKER_APT_MIRROR}"
DOCKER_KEYRING_PATH="${DOCKER_KEYRING_PATH:-$DEFAULT_KEYRING_PATH}"
LOCAL_DOCKER_GPG="${LOCAL_DOCKER_GPG:-/root/wdd/docker.gpg}"
LOG_LEVEL="${LOG_LEVEL:-INFO}"

# ==============================================================================
# Function Call Graph (ASCII)
# ==============================================================================
# main
#  |
#  +--> init_traps
#  |
#  +--> check_platform
#  |
#  +--> ensure_prerequisites
#  |
#  +--> detect_public_network
#  |      |
#  |      +--> can_fetch_url_head
#  |
#  +--> ensure_docker_gpg_key
#  |      |
#  |      +--> install_key_from_online
#  |      |      |
#  |      |      +--> require_cmd (curl, gpg)
#  |      |
#  |      +--> install_key_from_local
#  |
#  +--> configure_docker_repo
#  |
#  +--> install_docker_packages
#  |      |
#  |      +--> resolve_docker_version
#  |
#  +--> pin_docker_packages
#  |
#  +--> enable_docker_service
# ==============================================================================

# ==============================================================================
# Logging
# ==============================================================================

### Map log level string to numeric value.
### @param level_str string Level string (DEBUG/INFO/WARN/ERROR)
### @return 0 Always returns 0; outputs numeric level to stdout
### @require none
log_level_to_num() {
  case "${1:-INFO}" in
    DEBUG) echo 10 ;;
    INFO)  echo 20 ;;
    WARN)  echo 30 ;;
    ERROR) echo 40 ;;
    *)     echo 20 ;;
  esac
}

### Unified logger with level gating.
### @param level string Log level
### @param message string Message
### @return 0 Always returns 0
### @require date
log() {
  local level="${1:?level required}"
  shift
  local message="${*:-}"
  local now
  now="$(date '+%F %T')"

  local current_level_num wanted_level_num
  current_level_num="$(log_level_to_num "$LOG_LEVEL")"
  wanted_level_num="$(log_level_to_num "$level")"

  if [ "$wanted_level_num" -lt "$current_level_num" ]; then
    return 0
  fi

  # > Keep format stable for parsing by log collectors
  printf '%s [%s] %s: %s\n' "$now" "$level" "$SCRIPT_NAME" "$message" >&2
}

# ==============================================================================
# Error Handling & Traps
# ==============================================================================

### Trap handler for unexpected errors.
### @param exit_code int Exit code from failing command
### @return 0 Always returns 0
### @require none
on_error() {
  local exit_code="${1:-$EC_GENERAL}"
  log ERROR "Unhandled error occurred (exit_code=${exit_code})."
  exit "$exit_code"
}

### Trap handler for script exit.
### @param exit_code int Exit code
### @return 0 Always returns 0
### @require none
on_exit() {
  local exit_code="${1:-$EC_OK}"
  if [ "$exit_code" -eq 0 ]; then
    log INFO "Done."
  else
    log WARN "Exited with code ${exit_code}."
  fi
  return 0
}

### Initialize traps (ERR/INT/TERM/EXIT).
### @return 0 Success
### @require none
init_traps() {
  trap 'on_error $?' ERR
  trap 'log WARN "Interrupted (SIGINT)"; exit 130' INT
  trap 'log WARN "Terminated (SIGTERM)"; exit 143' TERM
  trap 'on_exit $?' EXIT
}

# ==============================================================================
# Privilege Helpers
# ==============================================================================

### Run a command as root (uses sudo if not root).
### @param cmd string Command to run
### @return 0 Success; non-zero on failure
### @require sudo (if not root)
run_root() {
  if [ "$(id -u)" -eq 0 ]; then
    # shellcheck disable=SC2068
    "$@"
  else
    # shellcheck disable=SC2068
    sudo "$@"
  fi
}

# ==============================================================================
# Dependency Checks
# ==============================================================================

### Ensure a command exists in PATH.
### @param cmd_name string Command name
### @return 0 If exists; 1 otherwise
### @require none
require_cmd() {
  local cmd_name="${1:?cmd required}"
  if ! command -v "$cmd_name" >/dev/null 2>&1; then
    log ERROR "Missing dependency: ${cmd_name}"
    return 1
  fi
  return 0
}

# ==============================================================================
# Platform Check
# ==============================================================================

### Check OS is Ubuntu and supported versions.
### @return 0 Supported; exits otherwise
### @require lsb_release, awk
check_platform() {
  require_cmd lsb_release || exit "$EC_DEPENDENCY"

  local distro version
  distro="$(lsb_release -is 2>/dev/null || true)"
  version="$(lsb_release -rs 2>/dev/null || true)"

  if [ "$distro" != "Ubuntu" ]; then
    log ERROR "Unsupported OS: ${distro}. This script supports Ubuntu only."
    exit "$EC_UNSUPPORTED_OS"
  fi

  case "$version" in
    18.04|20.04|22.04|24.04) ;;
    *)
      log ERROR "Unsupported Ubuntu version: ${version}. Supported: 18.04/20.04/22.04/24.04"
      exit "$EC_UNSUPPORTED_OS"
      ;;
  esac

  log INFO "Platform OK: ${distro} ${version}"
}

# ==============================================================================
# APT Prerequisites
# ==============================================================================

### Install required packages for repository/key management and Docker installation.
### @return 0 Success; exits on apt failures
### @require apt-get
ensure_prerequisites() {
  require_cmd apt-get || exit "$EC_DEPENDENCY"

  log INFO "Installing prerequisites (does NOT modify APT sources or proxy)..."
  # > apt update must work via your existing proxy+mirror scripts
  if ! run_root apt-get update; then
    log ERROR "apt-get update failed. Check APT proxy / mirror configuration."
    exit "$EC_APT_FAILURE"
  fi

  # > Keep dependencies minimal; curl/gpg used only for online key fetch.
  if ! run_root apt-get install -y ca-certificates gnupg lsb-release; then
    log ERROR "Failed to install prerequisites."
    exit "$EC_APT_FAILURE"
  fi

  log INFO "Prerequisites installed."
}

# ==============================================================================
# Public Network Reachability
# ==============================================================================

### Check whether we can fetch HTTP headers from a URL (lightweight reachability).
### @param test_url string URL to test
### @return 0 Reachable; 1 otherwise
### @require curl (optional; if missing returns 1)
can_fetch_url_head() {
  local test_url="${1:?url required}"

  if ! command -v curl >/dev/null 2>&1; then
    log WARN "curl not found; cannot test public network reachability via HTTP."
    return 1
  fi

  # > Use short timeout to avoid hanging in restricted networks
  curl -fsSI --max-time 3 "$test_url" >/dev/null 2>&1
}

### Detect whether public network access is available for Docker key fetch.
### @return 0 Online; 1 Offline/Uncertain
### @require none
detect_public_network() {
  local test_url="${DOCKER_APT_MIRROR%/}/gpg"

  log INFO "Detecting public network reachability: HEAD ${test_url}"
  if can_fetch_url_head "$test_url"; then
    log INFO "Public network reachable for Docker mirror."
    return 0
  fi

  log WARN "Public network NOT reachable (or curl missing). Will try local GPG key."
  return 1
}

# ==============================================================================
# Docker GPG Key Management
# ==============================================================================

### Install Docker GPG key from online source (mirror).
### @param gpg_url string GPG URL
### @param keyring_path string Keyring output path
### @return 0 Success; non-zero on failure
### @require curl, gpg, install, mkdir, chmod
install_key_from_online() {
  local gpg_url="${1:?gpg_url required}"
  local keyring_path="${2:?keyring_path required}"

  require_cmd curl || return 1
  require_cmd gpg  || return 1

  # > Write to temp then atomically install to avoid partial files
  local tmp_dir tmp_gpg
  tmp_dir="$(mktemp -d)"
  tmp_gpg="${tmp_dir}/docker.gpg"

  log INFO "Fetching Docker GPG key online: ${gpg_url}"
  curl -fsSL --max-time 10 "$gpg_url" | gpg --dearmor -o "$tmp_gpg"

  run_root mkdir -p "$(dirname "$keyring_path")"
  run_root install -m 0644 "$tmp_gpg" "$keyring_path"
  run_root chmod a+r "$keyring_path" || true

  rm -rf "$tmp_dir"
  log INFO "Docker GPG key installed: ${keyring_path}"
  return 0
}

### Install Docker GPG key from local file (offline-friendly).
### @param local_gpg_path string Local GPG file path
### @param keyring_path string Keyring output path
### @return 0 Success; 1 if local key missing; non-zero on other failures
### @require install, mkdir, chmod
install_key_from_local() {
  local local_gpg_path="${1:?local_gpg_path required}"
  local keyring_path="${2:?keyring_path required}"

  if [ ! -f "$local_gpg_path" ]; then
    log WARN "Local Docker GPG key not found: ${local_gpg_path}"
    return 1
  fi

  run_root mkdir -p "$(dirname "$keyring_path")"
  run_root install -m 0644 "$local_gpg_path" "$keyring_path"
  run_root chmod a+r "$keyring_path" || true

  log INFO "Docker GPG key installed from local: ${local_gpg_path} -> ${keyring_path}"
  return 0
}

### Ensure Docker GPG key exists, using online if reachable; otherwise local-only.
### Offline policy: if local key missing -> DO NOT proceed (exit).
### @param is_online int 0 online; 1 offline
### @return 0 Success; exits with EC_OFFLINE_NO_KEY when offline and no local key
### @require none
ensure_docker_gpg_key() {
  local is_online="${1:?is_online required}"

  # > If keyring already exists, reuse it (idempotent)
  if [ -f "$DOCKER_KEYRING_PATH" ]; then
    log INFO "Docker keyring already exists: ${DOCKER_KEYRING_PATH}"
    run_root chmod a+r "$DOCKER_KEYRING_PATH" || true
    return 0
  fi

  # > Determine local key candidate paths (priority order)
  local script_dir local_candidate
  script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

  if [ -n "$LOCAL_DOCKER_GPG" ]; then
    local_candidate="$LOCAL_DOCKER_GPG"
  elif [ -f "${script_dir}/docker.gpg" ]; then
    local_candidate="${script_dir}/docker.gpg"
  else
    local_candidate=""
  fi

  local gpg_url
  gpg_url="${DOCKER_APT_MIRROR%/}/gpg"

  if [ "$is_online" -eq 0 ]; then
    # Online: try online key fetch first; if fails, fallback to local if present.
    log DEBUG "Online mode: attempt online key install, fallback to local."
    if install_key_from_online "$gpg_url" "$DOCKER_KEYRING_PATH"; then
      return 0
    fi

    if [ -n "$local_candidate" ] && install_key_from_local "$local_candidate" "$DOCKER_KEYRING_PATH"; then
      return 0
    fi

    log ERROR "Failed to install Docker GPG key (online fetch failed and no usable local key)."
    exit "$EC_DEPENDENCY"
  fi

  # Offline: strictly local only; if missing -> do not proceed
  log INFO "Offline mode: install Docker GPG key from local only."
  if [ -n "$local_candidate" ] && install_key_from_local "$local_candidate" "$DOCKER_KEYRING_PATH"; then
    return 0
  fi

  log ERROR "Offline and local Docker GPG key is missing. Will NOT proceed (per policy)."
  exit "$EC_OFFLINE_NO_KEY"
}

# ==============================================================================
# Docker Repo Configuration
# ==============================================================================

### Configure Docker APT repository list file.
### @return 0 Success; exits on apt update failures
### @require dpkg, lsb_release, tee, apt-get
configure_docker_repo() {
  require_cmd dpkg        || exit "$EC_DEPENDENCY"
  require_cmd lsb_release  || exit "$EC_DEPENDENCY"
  require_cmd tee         || exit "$EC_DEPENDENCY"

  local codename arch list_file
  codename="$(lsb_release -cs)"
  arch="$(dpkg --print-architecture)"
  list_file="/etc/apt/sources.list.d/docker.list"

  log INFO "Configuring Docker APT repo: ${DOCKER_APT_MIRROR} (${codename}, ${arch})"
  # > Only touch docker repo; do not touch system sources.list
  run_root tee "$list_file" >/dev/null <<EOF
deb [arch=${arch} signed-by=${DOCKER_KEYRING_PATH}] ${DOCKER_APT_MIRROR} ${codename} stable
EOF

  if ! run_root apt-get update; then
    log ERROR "apt-get update failed after configuring Docker repo."
    exit "$EC_APT_FAILURE"
  fi

  log INFO "Docker APT repo configured: ${list_file}"
}

# ==============================================================================
# Docker Installation
# ==============================================================================

### Resolve Docker package version string from APT cache.
### @param docker_version string Desired version ("20.10" or "20.10.15")
### @return 0 Success and echoes full apt version string; exits if not found
### @require apt-cache, awk, grep, sort, head
resolve_docker_version() {
  local docker_version="${1:?docker_version required}"

  require_cmd apt-cache || exit "$EC_DEPENDENCY"
  require_cmd awk       || exit "$EC_DEPENDENCY"
  require_cmd grep      || exit "$EC_DEPENDENCY"
  require_cmd sort      || exit "$EC_DEPENDENCY"
  require_cmd head      || exit "$EC_DEPENDENCY"

  local resolved=""
  # > apt-cache madison output includes epoch, keep it for apt-get install
  if [[ "$docker_version" =~ ^[0-9]+\.[0-9]+$ ]]; then
    # Pick newest patch/build for that major.minor
    resolved="$(
      apt-cache madison docker-ce \
        | awk -F'|' '{gsub(/ /,"",$2); print $2}' \
        | grep -E "^[0-9]+:${docker_version}([.-]|\~)" \
        | sort -rV \
        | head -1 || true
    )"
  elif [[ "$docker_version" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    resolved="$(
      apt-cache madison docker-ce \
        | awk -F'|' '{gsub(/ /,"",$2); print $2}' \
        | grep -E "^[0-9]+:${docker_version}.*" \
        | head -1 || true
    )"
  else
    log ERROR "Invalid DOCKER_VERSION format: ${docker_version} (expect 20.10 or 20.10.15)"
    exit "$EC_GENERAL"
  fi

  if [ -z "$resolved" ]; then
    log ERROR "Cannot find Docker version '${docker_version}' from APT. Check repo/mirror and apt proxy."
    exit "$EC_APT_FAILURE"
  fi

  echo "$resolved"
  return 0
}

### Install Docker packages via APT.
### @return 0 Success; exits on failure
### @require apt-get, systemctl
install_docker_packages() {
  require_cmd apt-get || exit "$EC_DEPENDENCY"

  local full_version
  full_version="$(resolve_docker_version "$DOCKER_VERSION")"

  log INFO "Installing Docker packages: docker-ce=${full_version}"
  # > Compose: use docker-compose-plugin (no curl downloading binaries)
  if ! run_root apt-get install -y \
      "docker-ce=${full_version}" \
      "docker-ce-cli=${full_version}" \
      "docker-ce-rootless-extras=${full_version}" \
      containerd.io \
      docker-buildx-plugin \
      docker-compose-plugin; then
    log ERROR "Docker installation failed."
    exit "$EC_APT_FAILURE"
  fi

  # > Optional: provide docker-compose legacy command compatibility
  if ! command -v docker-compose >/dev/null 2>&1; then
    if [ -x /usr/libexec/docker/cli-plugins/docker-compose ]; then
      run_root ln -sf /usr/libexec/docker/cli-plugins/docker-compose /usr/local/bin/docker-compose || true
    fi
  fi

  log INFO "Docker packages installed."
}

### Pin Docker packages to avoid unintended upgrades.
### @return 0 Success; non-zero on failures (non-fatal)
### @require apt-mark
pin_docker_packages() {
  if ! command -v apt-mark >/dev/null 2>&1; then
    log WARN "apt-mark not found; skip pinning."
    return 0
  fi

  log INFO "Holding Docker packages (prevent auto-upgrade)..."
  run_root apt-mark hold \
    docker-ce docker-ce-cli docker-ce-rootless-extras containerd.io \
    docker-buildx-plugin docker-compose-plugin >/dev/null 2>&1 || true
  return 0
}

### Enable and start Docker service, then verify versions.
### @return 0 Success; exits on failure to enable docker
### @require systemctl, docker
enable_docker_service() {
  require_cmd systemctl || exit "$EC_DEPENDENCY"

  log INFO "Enabling and starting docker service..."
  run_root systemctl enable --now docker

  # > Verification should not hard-fail the whole script
  if command -v docker >/dev/null 2>&1; then
    docker --version || true
    docker compose version || true
  fi
  if command -v docker-compose >/dev/null 2>&1; then
    docker-compose --version || true
  fi

  log INFO "Docker service enabled."
}

# ==============================================================================
# Main
# ==============================================================================

### Main entrypoint.
### @return 0 Success; non-zero on failure
### @require none
main() {
  init_traps
  log INFO "Starting Docker installer (v${SCRIPT_VERSION})..."

  check_platform
  ensure_prerequisites

  local is_online=1
  if detect_public_network; then
    is_online=0
  fi

  ensure_docker_gpg_key "$is_online"
  configure_docker_repo
  install_docker_packages
  pin_docker_packages
  enable_docker_service

  log INFO "All tasks completed successfully."
  exit "$EC_OK"
}

main "$@"