#!/bin/bash

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

# 生成默认的 config.toml 配置文件（覆盖可能存在的失败配置）
# echo "生成默认的 config.toml 配置..."
# sudo containerd config default | sudo tee /etc/containerd/config.toml > /dev/null
# sudo containerd config default | sed 's/version = "3"/version = "2"/' | sudo tee /etc/containerd/config.toml > /dev/null

#
CONFIG_PATH="/etc/containerd/certs.d"
NEW_STORAGE_PATH="/data/containerd"
OLD_STORAGE_PATH="/var/lib/containerd"
OLD_STATE_PATH="/run/containerd"  # 旧的 state 目录

# 查找 [plugins.'io.containerd.grpc.v1.cri'.registry] 的位置并修改或添加 config_path 配置
echo '配置 containerd 的 config.toml...'

line_number=$(grep -n -E '^\s*\[plugins.(\"|\x27)io.containerd.grpc.v1.cri(\"|\x27).registry\]' /etc/containerd/config.toml | cut -d':' -f1)

if [ -n "$line_number" ]; then
  # 检查下一行是否包含 `config_path`
  next_line=$((line_number + 1))
  if grep -q "^\s*config_path\s*=" /etc/containerd/config.toml; then
    # 若找到 `config_path`，则直接替换为新的值
    sudo sed -i "${next_line}s|config_path\s*=.*|config_path = '$CONFIG_PATH'|" /etc/containerd/config.toml
    echo "已修改 config_path 为 $CONFIG_PATH。"
  else
    echo "未找到 config_path，未执行任何修改。"
  fi
else
  echo "未找到 [plugins.'io.containerd.grpc.v1.cri'.registry] 段落，请检查 config.toml 文件格式。"
fi

# 检查配置文件
sudo containerd config dump

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

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

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

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

# 删除所有的旧的镜像
sudo nerdctl --namespace k8s.io image prune -a --force
sudo nerdctl --namespace default image prune -a --force

# 获取迁移后的新目录大小
initial_size=$(sudo du -sb "$NEW_STORAGE_PATH" | awk '{print $1}')

# 配置加速镜像的主机和路径
declare -A mirrors
mirrors=(
  ["docker.io"]="https://docker.io https://docker.unsee.tech https://dockerhub.icu"
  ["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"
  sudo mkdir -p "$DIR"
  sudo 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

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

# 设置 containerd.sock 的权限
echo "设置 containerd.sock 的访问权限..."
# sudo mkdir -p /run/containerd
# sudo touch /run/containerd/containerd.sock
sudo chown root:containerd /run/containerd/containerd.sock
sudo chmod 660 /run/containerd/containerd.sock

# 将当前用户添加到 containerd 组，以允许使用 ctr
echo "将用户 $CURRENT_USER 添加到 containerd 组..."
sudo groupadd -f containerd  # 确保 containerd 组存在
sudo usermod -aG containerd "$CURRENT_USER"

# 提示用户重新登录以应用组权限更改
echo "完成。请重新登录会话以应用对 group 的更改，使 $CURRENT_USER 可以使用 ctr。"

# 测试配置是否生效并获取新的目录大小
echo '测试配置是否生效...'
if sudo ctr --namespace=default image pull --hosts-dir /etc/containerd/certs.d docker.io/library/alpine:latest; then
  echo "镜像加速配置成功！"
  
  # 检查拉取后目录大小
  final_size=$(sudo du -sb "$NEW_STORAGE_PATH" | awk '{print $1}')
  
  if [ "$final_size" -gt "$initial_size" ]; then
    echo "存储路径迁移成功，镜像数据已写入新路径。"
  else
    echo "存储路径迁移失败，新路径未见数据变化，请检查配置。"
  fi
else
  echo "镜像加速配置失败，请检查配置。"
fi