#!/usr/bin/env bash
set -euo pipefail

# 教学用：在中国地区自动安装 Docker Engine
# 功能：
# 1) 识别系统类型 (Debian/Ubuntu, RHEL/CentOS/Alma/Rocky, Fedora)
# 2) 自动挑选国内可用且尽量更快的 Docker 下载镜像
# 3) 通过官方 get.docker.com 脚本安装（并传入镜像环境变量）
# 4) 可选：为 Docker Hub 配置一个国内 registry mirror（若 daemon.json 不存在则写入）
# 5) 启动并开机自启，当前用户加入 docker 组

if [[ ${EUID:-$(id -u)} -ne 0 ]]; then SUDO="sudo"; else SUDO=""; fi

# 可选参数：
#   --no-registry-mirror   跳过为 Docker Hub 写入 registry-mirrors
#   --force-mirror         即使已有 daemon.json 也强制写入/更新 registry-mirrors（优先 jq 合并，不在则备份后覆盖）
#   --mirror "URL1,URL2"    自定义镜像加速器列表（逗号或空格分隔）
#   --max-attempts N       下载镜像尝试次数（默认 5）

DISABLE_MIRROR=0
FORCE_MIRROR=0
CUSTOM_MIRRORS=""
MAX_ATTEMPTS=5

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-registry-mirror) DISABLE_MIRROR=1; shift ;;
    --force-mirror)       FORCE_MIRROR=1; shift ;;
    --mirror)             shift; CUSTOM_MIRRORS="${1:-}"; shift ;;
    --max-attempts)       shift; MAX_ATTEMPTS=${1:-5}; shift ;;
    -h|--help)
      cat <<USAGE
用法: $0 [选项]
  --no-registry-mirror   跳过为 Docker Hub 写入 registry-mirrors
  --force-mirror         强制写入/更新 registry-mirrors（有 jq 则合并，无 jq 则备份后覆盖）
  --mirror "URL1,URL2"    指定一个或多个加速器，逗号/空格分隔（默认自动测速选择）
  --max-attempts N       下载镜像尝试次数（默认 5）
  -h, --help             显示帮助
USAGE
      exit 0
      ;;
    *) warn "忽略未知参数：$1"; shift ;;
  esac
done

log()  { printf "[INFO] %s\n" "$*"; }
warn() { printf "[WARN] %s\n" "$*" 1>&2; }
err()  { printf "[ERR ] %s\n"  "$*" 1>&2; }

need_cmd() { command -v "$1" >/dev/null 2>&1; }

install_prereqs() {
  if need_cmd curl && need_cmd awk; then return 0; fi
  if command -v apt-get >/dev/null 2>&1; then
    $SUDO apt-get update -y || true
    $SUDO apt-get install -y curl ca-certificates gnupg lsb-release awk || true
  elif command -v dnf >/dev/null 2>&1; then
    $SUDO dnf install -y curl ca-certificates gawk || true
  elif command -v yum >/dev/null 2>&1; then
    $SUDO yum install -y curl ca-certificates gawk || true
  fi
  if ! need_cmd curl; then err "缺少 curl，请手动安装后重试"; exit 1; fi
  if ! need_cmd awk;  then err "缺少 awk，请手动安装后重试";  exit 1; fi
}

detect_family() {
  local id like
  id=""; like=""
  if [[ -r /etc/os-release ]]; then
    . /etc/os-release
    id=${ID:-}
    like=${ID_LIKE:-}
  fi
  case "$id" in
    ubuntu|debian|raspbian) echo debian; return;;
    fedora) echo fedora; return;;
    centos|rhel|rocky|almalinux|ol|anolis) echo rhel; return;;
  esac
  if [[ "$like" =~ (debian) ]]; then echo debian; return; fi
  if [[ "$like" =~ (rhel|centos|fedora) ]]; then echo rhel; return; fi
  echo unknown
}

pick_fastest_download_base() {
  # 返回：Name|DOWNLOAD_URL  （REPO_URL = ${DOWNLOAD_URL}/linux）
  # 支持通过环境变量 EXCLUDE_URLS（用空格分隔）来排除已失败的地址
  local -a candidates=(
    "Aliyun|https://mirrors.aliyun.com/docker-ce"
    "Tsinghua|https://mirrors.tuna.tsinghua.edu.cn/docker-ce"
    "USTC|https://mirrors.ustc.edu.cn/docker-ce"
    "Huawei|https://repo.huaweicloud.com/docker-ce"
    "Official|https://download.docker.com"
  )
  local best_name="" best_url="" best_time="999999"
  local item name url code time http_and_time skip
  for item in "${candidates[@]}"; do
    name="${item%%|*}"; url="${item#*|}"
    skip=0
    for ex in ${EXCLUDE_URLS:-}; do [[ "$url" == "$ex" ]] && skip=1 && break; done
    [[ $skip -eq 1 ]] && continue
    http_and_time=$(curl -m 4 -L -s -o /dev/null -w "%{http_code} %{time_total}" "${url}/linux/") || true
    code="${http_and_time%% *}"; time="${http_and_time#* }"
    [[ -z "$time" ]] && time="999999"
    if [[ "$code" =~ ^2 ]] || [[ "$code" =~ ^3 ]]; then
      log "镜像可用：$name $url (${time}s)"
      if awk -v t="$time" -v b="$best_time" 'BEGIN{exit (t < b) ? 0 : 1}'; then
        best_time="$time"; best_url="$url"; best_name="$name"
      fi
    else
      warn "镜像不可用：$name $url (HTTP $code)"
    fi
  done
  if [[ -z "$best_url" ]]; then
    warn "未探测到可用镜像，使用官方：https://download.docker.com"
    best_name="Official"; best_url="https://download.docker.com"
  fi
  printf "%s|%s\n" "$best_name" "$best_url"
}

pick_fastest_registry_mirror() {
  # 返回：Name|URL ；挑选 Docker Hub 加速器
  local -a mirrors=(
    "Tencent |https://mirror.ccs.tencentyun.com"
    "USTC    |https://docker.mirrors.ustc.edu.cn"
    "1PanelP |https://dockerproxy.1panel.live"
    "1Panel  |https://docker.1panel.live"
    "1PanelX |https://docker.1panelproxy.com"
    "1Proxy  |https://proxy.1panel.live"
    "NetEase |https://hub-mirror.c.163.com"
    "Proxy   |https://dockerproxy.com"
  )
  local best_name="" best_url="" best_time="999999"
  local item name url code time http_and_time
  for item in "${mirrors[@]}"; do
    name="${item%%|*}"; url="${item#*|}"
    http_and_time=$(curl -m 4 -L -s -o /dev/null -w "%{http_code} %{time_total}" "${url}/v2/") || true
    code="${http_and_time%% *}"; time="${http_and_time#* }"
    [[ -z "$time" ]] && time="999999"
    if [[ "$code" =~ ^2 ]] || [[ "$code" =~ ^3 ]] || [[ "$code" =~ ^4 ]]; then
      log "Hub 加速器可用：$name $url (${time}s)"
      if awk -v t="$time" -v b="$best_time" 'BEGIN{exit (t < b) ? 0 : 1}'; then
        best_time="$time"; best_url="$url"; best_name="$name"
      fi
    else
      warn "Hub 加速器不可用：$name $url (HTTP $code)"
    fi
  done
  if [[ -z "$best_url" ]]; then
    warn "未探测到可用 Hub 加速器，将不写入 registry-mirrors"
    return 1
  fi
  printf "%s|%s\n" "$best_name" "$best_url"
}

# 返回按速度排序的可用加速器（空格分隔）
pick_registry_mirrors_sorted() {
  local -a urls=(
    "https://mirror.ccs.tencentyun.com"
    "https://docker.mirrors.ustc.edu.cn"
    "https://dockerproxy.1panel.live"
    "https://docker.1panel.live"
    "https://docker.1panelproxy.com"
    "https://proxy.1panel.live"
    "https://hub-mirror.c.163.com"
    "https://dockerproxy.com"
  )
  local lines="" u http_and_time code time
  for u in "${urls[@]}"; do
    http_and_time=$(curl -m 4 -L -s -o /dev/null -w "%{http_code} %{time_total}" "$u/v2/") || true
    code="${http_and_time%% *}"; time="${http_and_time#* }"
    [[ -z "$time" ]] && time="999999"
    if [[ "$code" =~ ^2 ]] || [[ "$code" =~ ^3 ]] || [[ "$code" =~ ^4 ]]; then
      log "Hub 加速器可用：$u (${time}s)"
      lines+="$time $u\n"
    else
      warn "Hub 加速器不可用：$u (HTTP $code)"
    fi
  done
  if [[ -z "$lines" ]]; then return 1; fi
  printf "%s" "$lines" | LC_ALL=C sort -n | awk '{print $2}' | tr '\n' ' '
}

mirrors_to_json_array() {
  # 输入：以空格分隔的镜像 URL 列表；输出：JSON 数组字符串
  # 例："https://a https://b" -> ["https://a","https://b"]
  local arr=($1)
  local out="["
  local i
  for ((i=0;i<${#arr[@]};i++)); do
    out+="\"${arr[$i]}\""
    if (( i < ${#arr[@]}-1 )); then out+=","; fi
  done
  out+="]"
  printf "%s" "$out"
}

ensure_service_started() {
  if command -v systemctl >/dev/null 2>&1; then
    $SUDO systemctl daemon-reload || true
    $SUDO systemctl enable --now docker || $SUDO systemctl start docker || true
  elif command -v service >/dev/null 2>&1; then
    $SUDO service docker start || true
  else
    warn "未检测到 systemd/service，尝试直接后台启动 dockerd"
    if need_cmd dockerd; then nohup $SUDO dockerd >/var/log/dockerd.log 2>&1 & fi
  fi
}

post_install() {
  local target_user
  if [[ -n "${SUDO_USER:-}" && "${SUDO_USER:-}" != "root" ]]; then
    target_user="$SUDO_USER"
  else
    target_user="$(id -un)"
  fi
  $SUDO groupadd -f docker || true
  if [[ "$target_user" != "root" ]]; then
    $SUDO usermod -aG docker "$target_user" || true
    log "已将用户 $target_user 加入 docker 组（重新登录后生效）"
  fi
}

main() {
  install_prereqs

  local family
  family=$(detect_family)
  if [[ "$family" == "unknown" ]]; then
    warn "未识别的系统，仍尝试安装（若失败请改为手动仓库方式）"
  else
    log "检测到系统族：$family"
  fi

  local chosen name base_url
  local tmp_script
  tmp_script=$(mktemp /tmp/get-docker.XXXXXX.sh)
  log "下载官方安装脚本 get.docker.com ..."
  if ! curl -fsSL -m 30 https://get.docker.com -o "$tmp_script"; then
    err "下载 get.docker.com 失败，请检查网络或稍后再试"
    exit 2
  fi

  # 循环尝试：每次挑选当前最快镜像，失败则加入排除列表并继续
  local attempt=1 success=0
  EXCLUDE_URLS=""
  while [[ $attempt -le $MAX_ATTEMPTS ]]; do
    chosen=$(pick_fastest_download_base)
    name="${chosen%%|*}"; base_url="${chosen#*|}"
    log "第 ${attempt}/${MAX_ATTEMPTS} 次安装尝试：$name -> $base_url"

    if DOWNLOAD_URL="$base_url" REPO_URL="$base_url/linux" sh "$tmp_script"; then
      success=1
      break
    else
      warn "使用镜像 $base_url 安装失败，切换下一个镜像重试"
      EXCLUDE_URLS="$EXCLUDE_URLS $base_url"
      attempt=$((attempt+1))
    fi
  done
  if [[ $success -ne 1 ]]; then
    err "所有镜像尝试均失败，请检查网络或改用手动仓库方法"
    exit 3
  fi

  # 配置 registry mirror（支持自定义/强制覆盖/跳过）
  if [[ $DISABLE_MIRROR -eq 0 ]]; then
    local mirror_list=""
    if [[ -n "$CUSTOM_MIRRORS" ]]; then
      mirror_list="$(echo "$CUSTOM_MIRRORS" | tr ',' ' ')"
      log "使用自定义加速器：$mirror_list"
    else
      if mirror_list=$(pick_registry_mirrors_sorted); then
        log "写入测速排序后的加速器列表：$mirror_list"
      else
        warn "未找到可用加速器，跳过写入 registry-mirrors"
      fi
    fi

    if [[ -n "$mirror_list" ]]; then
      local json_arr
      json_arr=$(mirrors_to_json_array "$mirror_list")
      $SUDO install -d -m 0755 /etc/docker

      if [[ -f /etc/docker/daemon.json ]]; then
        if command -v jq >/dev/null 2>&1; then
          local tmpf
          tmpf=$(mktemp)
          if $SUDO jq --argjson arr "$json_arr" '."registry-mirrors"=$arr' /etc/docker/daemon.json > "$tmpf" 2>/dev/null; then
            $SUDO mv "$tmpf" /etc/docker/daemon.json
            log "已用 jq 更新 /etc/docker/daemon.json 的 registry-mirrors"
          else
            warn "jq 更新失败，跳过更新（可使用 --force-mirror 覆盖）"
            rm -f "$tmpf" 2>/dev/null || true
          fi
        elif [[ $FORCE_MIRROR -eq 1 ]]; then
          local bak="/etc/docker/daemon.json.bak-$(date +%Y%m%d%H%M%S)"
          $SUDO cp /etc/docker/daemon.json "$bak"
          warn "未检测到 jq，已备份旧配置到 $bak，并覆盖写入 registry-mirrors"
          $SUDO sh -c "cat > /etc/docker/daemon.json" <<EOF
{
  "registry-mirrors": $json_arr
}
EOF
        else
          warn "/etc/docker/daemon.json 已存在，为避免破坏现有配置未写入。可用 --force-mirror 或预装 jq 后重试。"
        fi
      else
        $SUDO tee /etc/docker/daemon.json >/dev/null <<EOF
{
  "registry-mirrors": $json_arr
}
EOF
      fi
    fi
  else
    log "按参数要求跳过 registry-mirrors 配置"
  fi

  ensure_service_started
  post_install

  log "安装完成！建议执行："
  printf "  docker version\n"
  printf "  docker run --rm hello-world\n"
}

main "$@"