#!/bin/bash

# 获取当前执行脚本的用户
CURRENT_USER=$(whoami)
echo "当前用户：$CURRENT_USER"

# 配置路径
CONFIG_PATH="$HOME/.config/containerd/certs.d"
CONFIG_FILE="$HOME/.config/containerd/config.toml"
NEW_STORAGE_PATH="/data/zly/containerd"
OLD_STORAGE_PATH="$HOME/.local/share/containerd"  # 旧的 rootless 存储路径
OLD_STATE_PATH="$XDG_RUNTIME_DIR/containerd"       # 旧的 state 路径

# 创建配置目录
mkdir -p "$CONFIG_PATH"
mkdir -p "$HOME/.config/containerd"

# 生成默认的 config.toml 配置文件（如果不存在）
if [ ! -f "$CONFIG_FILE" ]; then
  echo "生成默认的 config.toml 配置..."
  containerd config default | sed 's|/var/lib/containerd|'"$NEW_STORAGE_PATH"'|g' > "$CONFIG_FILE"
fi

# 配置 registry
echo '配置 containerd 的 config.toml...'

# 修改 registry 的 `config_path` 参数
line_number=$(grep -n -E '^\s*\[plugins.(\"|\x27)io.containerd.grpc.v1.cri(\"|\x27).registry\]' "$CONFIG_FILE" | cut -d':' -f1)

if [ -n "$line_number" ]; then
  next_line=$((line_number + 1))
  if grep -q "^\s*config_path\s*=" "$CONFIG_FILE"; then
    sed -i "${next_line}s|config_path\s*=.*|config_path = '$CONFIG_PATH'|" "$CONFIG_FILE"
    echo "已修改 config_path 为 $CONFIG_PATH。"
  else
    echo "未找到 config_path，未执行任何修改。"
  fi
else
  echo "未找到 [plugins.'io.containerd.grpc.v1.cri'.registry] 段落，请检查 config.toml 文件格式。"
fi

# 修改 root 和 state 路径
echo "更新 containerd 存储路径配置..."
sed -i "s|^\s*root\s*=.*|root = \"$NEW_STORAGE_PATH\"|" "$CONFIG_FILE"
sed -i "s|^\s*state\s*=.*|state = \"$NEW_STORAGE_PATH/state\"|" "$CONFIG_FILE"

# 创建新的存储目录和状态目录
echo "创建新的存储目录 $NEW_STORAGE_PATH 和状态目录..."
mkdir -p "$NEW_STORAGE_PATH"
mkdir -p "$NEW_STORAGE_PATH/state"
chmod -R 700 "$NEW_STORAGE_PATH"

# 迁移旧的存储数据到新的存储路径
if [ -d "$OLD_STORAGE_PATH" ]; then
  echo "迁移旧存储数据到新路径..."
  rsync -a "$OLD_STORAGE_PATH/" "$NEW_STORAGE_PATH/"
  
  if [ $? -eq 0 ]; then
    echo "存储数据迁移完成，删除旧的存储路径..."
    rm -rf "$OLD_STORAGE_PATH"
    echo "旧的存储路径已删除。"
  else
    echo "存储数据迁移失败，保留旧的存储路径。"
  fi
else
  echo "旧的存储路径 $OLD_STORAGE_PATH 不存在，无需迁移。"
fi

# 配置镜像源的主机和路径
declare -A mirrors
mirrors=(
  ["docker.io"]="https://docker.io https://docker.unsee.tech https://dockerhub.icu https://upnuemce.mirror.aliyuncs.com"
  ["registry.k8s.io"]="https://registry.k8s.io https://k8s.m.daocloud.io"
  ["docker.elastic.co"]="https://docker.elastic.co https://elastic.m.daocloud.io"
  ["gcr.io"]="https://gcr.io https://gcr.m.daocloud.io"
  ["ghcr.io"]="https://ghcr.io https://ghcr.m.daocloud.io"
  ["k8s.gcr.io"]="https://k8s.gcr.io https://k8s-gcr.m.daocloud.io"
  ["mcr.microsoft.com"]="https://mcr.microsoft.com https://mcr.m.daocloud.io"
  ["nvcr.io"]="https://nvcr.io https://nvcr.m.daocloud.io"
  ["quay.io"]="https://quay.io https://quay.m.daocloud.io"
  ["registry.jujucharms.com"]="https://registry.jujucharms.com https://jujucharms.m.daocloud.io"
  ["rocks.canonical.com"]="https://rocks.canonical.com https://rocks-canonical.m.daocloud.io"
)

# 创建 hosts.toml 配置文件
echo '正在配置镜像加速...'
for registry in "${!mirrors[@]}"; do
  IFS=' ' read -r server host1 host2 host3 <<<"${mirrors[$registry]}"
  DIR="$CONFIG_PATH/$registry"
  mkdir -p "$DIR"
  tee "$DIR/hosts.toml" > /dev/null <<EOF
server = "$server"

[host."$host1"]
  capabilities = ["pull", "resolve", "push"]

[host."$host2"]
  capabilities = ["pull", "resolve", "push"]

[host."$host3"]
  capabilities = ["pull", "resolve", "push"]
EOF
done

# 删除旧的镜像缓存
nerdctl --namespace k8s.io image prune -a --force
nerdctl --namespace default image prune -a --force

# 重启 containerd 服务
echo "重启 containerd 服务..."
systemctl --user daemon-reload
systemctl --user restart containerd
sudo systemctl daemon-reload
sudo systemctl restart containerd

# 测试配置是否生效
echo '测试配置是否生效...'
# nerdctl --namespace=default pull docker.io/bioconductor/cuda:devel
# nerdctl --namespace=default pull docker.io/bioconductor/cuda:devel-R-devel
# nerdctl --namespace=default --hosts-dir="$HOME/.config/containerd/certs.d" pull docker.io/ollama/ollama:latest
if nerdctl --namespace=default --hosts-dir="$CONFIG_PATH" pull docker.io/library/alpine:latest; then
  echo "镜像加速配置成功！"
else
  echo "镜像加速配置失败，请检查配置。"
fi