#!/bin/bash
set -euo pipefail

# 说明:
# 在 master、node1、node2 三台节点上分别拉取指定镜像, 并导入到 containerd (k8s.io 命名空间)
# 不通过主机分发镜像归档, 而是每台节点各自拉取/导入。
#
# 使用示例:
#   chmod +x k8s-image-pull-and-import.sh
#   ./k8s-image-pull-and-import.sh beppeb/devstar-controller-manager:3.0.0.without_istio
#
# 可选环境变量:
#   SSH_KEY  指定私钥路径 (默认: ~/.ssh/id_rsa, 若存在自动携带)

echo "==== K8s 镜像拉取并导入 containerd ===="

if [ $# -lt 1 ]; then
  echo "用法: $0 <IMAGE[:TAG]>"
  echo "示例: $0 beppeb/devstar-controller-manager:3.0.0.without_istio"
  exit 1
fi

IMAGE_INPUT="$1"

# 规范化镜像名, 若无 registry 前缀则补全 docker.io/
normalize_image() {
  local img="$1"
  if [[ "$img" != */*/* ]]; then
    # 只有一个斜杠(如 library/nginx 或 beppeb/devstar-...): 仍可能缺少 registry
    # Docker 的默认 registry 是 docker.io
    echo "docker.io/${img}"
  else
    echo "$img"
  fi
}

CANONICAL_IMAGE=$(normalize_image "$IMAGE_INPUT")
echo "目标镜像: ${CANONICAL_IMAGE}"

# 节点列表: 与 k8s-step1-prepare-env.sh 风格一致
NODES=("172.17.0.15:master" "172.17.0.43:node1" "172.17.0.34:node2")

# 本机 IP 与 SSH 选项
LOCAL_IP=$(ip route get 1 | awk '{print $7; exit}')
SSH_OPTS='-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o BatchMode=yes'
SSH_KEY_PATH=${SSH_KEY:-$HOME/.ssh/id_rsa}
[ -f "$SSH_KEY_PATH" ] && SSH_ID="-i $SSH_KEY_PATH" || SSH_ID=""

run_remote() {
  local ip="$1"; shift
  local cmd="$*"
  if [ "$ip" = "$LOCAL_IP" ]; then
    bash -lc "$cmd"
  else
    ssh $SSH_OPTS $SSH_ID ubuntu@"$ip" "$cmd"
  fi
}

# 在远端节点执行: 使用 docker 或 containerd 拉取镜像, 并确保导入到 containerd k8s.io
remote_pull_and_import_cmd() {
  local image="$1"
  # 注意: 使用单引号包裹, 传到远端后再展开变量
  cat <<'EOF_REMOTE'
set -euo pipefail

IMAGE_REMOTE="$IMAGE_PLACEHOLDER"

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

echo "[\"$(hostname)\"] 处理镜像: ${IMAGE_REMOTE}"

# 优先尝试 docker 拉取, 成功后直接导入 containerd (无需落盘)
if has_cmd docker; then
  echo "[\"$(hostname)\"] 使用 docker pull"
  sudo docker pull "${IMAGE_REMOTE}"
  echo "[\"$(hostname)\"] 导入到 containerd (k8s.io)"
  sudo docker save "${IMAGE_REMOTE}" | sudo ctr -n k8s.io images import - >/dev/null
else
  echo "[\"$(hostname)\"] 未检测到 docker, 尝试使用 containerd 拉取"
  # containerd 直接拉取到 k8s.io 命名空间
  sudo ctr -n k8s.io images pull --all-platforms "${IMAGE_REMOTE}"
fi

# 规范化 tag: 若镜像缺少 docker.io 前缀, 在 containerd 内补齐一份别名
NEED_PREFIX=0
if [[ "${IMAGE_REMOTE}" != docker.io/* ]]; then
  NEED_PREFIX=1
fi

if [ "$NEED_PREFIX" -eq 1 ]; then
  # 仅当不存在 docker.io/ 前缀时, 补一个 docker.io/ 的 tag, 方便与清单匹配
  # 计算补齐后的名字
  if [[ "${IMAGE_REMOTE}" == */*/* ]]; then
    # 已有显式 registry, 不重复打 tag
    :
  else
    FIXED="docker.io/${IMAGE_REMOTE}"
    echo "[\"$(hostname)\"] 为 containerd 打标签: ${FIXED}"
    sudo ctr -n k8s.io images tag "${IMAGE_REMOTE}" "${FIXED}" || true
  fi
fi

echo "[\"$(hostname)\"] 验证镜像是否存在于 containerd:"
sudo ctr -n k8s.io images ls | grep -E "$(printf '%s' "${IMAGE_REMOTE}" | sed 's/[\/.\-]/\\&/g')" || true
EOF_REMOTE
}

# 遍历节点执行
for node in "${NODES[@]}"; do
  IFS=':' read -r ip hostname <<< "$node"
  echo "==== 在 ${hostname} (${ip}) 执行镜像拉取与导入 ===="
  # 将占位符替换为实际镜像并远程执行
  remote_script=$(remote_pull_and_import_cmd "$CANONICAL_IMAGE")
  # 安全替换占位符为镜像名
  remote_script=${remote_script//\$IMAGE_PLACEHOLDER/$CANONICAL_IMAGE}
  run_remote "$ip" "$remote_script"
  echo ""
done

echo "==== 完成 ===="