#!/bin/bash

# ============================================
# AIOT Platform - 清理脚本
# 清理部署主机上的旧镜像 / 停止容器 / 构建缓存，释放存储空间
# ============================================

set -e

# ---- 默认参数 ----
KEEP=3
PRUNE_VOLUMES=false
REGISTRY_HOST="localhost:5000"

usage() {
    cat <<EOF
用法: $0 [options]

Options:
  --keep=N            每个服务保留最近 N 个本地镜像（默认 3）
  --prune-volumes     额外清理未使用的 Docker volume（默认不清，避免误删数据）
  --registry=HOST     本地 docker images 仓库前缀（默认 localhost:5000）
  -h, --help          帮助

示例：
  $0 --keep=3
  $0 --keep=2 --prune-volumes
EOF
}

for arg in "$@"; do
    case "$arg" in
        --keep=*)         KEEP="${arg#*=}" ;;
        --prune-volumes)  PRUNE_VOLUMES=true ;;
        --registry=*)     REGISTRY_HOST="${arg#*=}" ;;
        -h|--help)        usage; exit 0 ;;
        *)                echo "未知参数: $arg"; usage; exit 1 ;;
    esac
done

if ! [[ "$KEEP" =~ ^[0-9]+$ ]] || [ "$KEEP" -lt 1 ]; then
    echo "❌ --keep 必须为正整数"
    exit 1
fi

# 颜色输出
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }

log_info "========================================="
log_info "AIOT Platform 清理开始 (keep=${KEEP})"
log_info "========================================="

log_info "当前磁盘使用情况:"
df -h | grep -E "Filesystem|/$" || true
echo ""
log_info "当前 Docker 磁盘使用:"
docker system df
echo ""

log_info "清理停止的容器..."
docker container prune -f

log_info "清理悬空镜像..."
docker image prune -f

log_info "清理旧版本镜像（每个服务保留最近 ${KEEP} 个）..."

SERVICES=(
    "viewsh-gateway"
    "viewsh-module-system-server"
    "viewsh-module-infra-server"
    "viewsh-module-iot-server"
    "viewsh-module-iot-gateway"
    "viewsh-module-ops-server"
)

for service in "${SERVICES[@]}"; do
    log_info "处理服务: ${service}"

    # 按创建时间降序取 ID 列表，跳过 latest tag，保留前 KEEP 个
    mapfile -t ids_to_delete < <(
        docker images "${REGISTRY_HOST}/${service}" \
            --format '{{.CreatedAt}}|{{.ID}}|{{.Tag}}' \
            | grep -v '|latest$' \
            | sort -r \
            | awk -F'|' -v k="$KEEP" 'NR > k {print $2}'
    )

    if [ "${#ids_to_delete[@]}" -eq 0 ]; then
        log_info "  └─ 无可清理镜像"
        continue
    fi

    log_info "  └─ 删除 ${#ids_to_delete[@]} 个旧镜像"
    # 去重后批量删
    printf '%s\n' "${ids_to_delete[@]}" | sort -u | xargs -r docker rmi -f 2>/dev/null || true
done

if [ "$PRUNE_VOLUMES" = true ]; then
    log_warn "清理未使用的 volume（--prune-volumes 已启用）"
    docker volume prune -f
else
    log_info "跳过 volume 清理（如需清理请加 --prune-volumes）"
fi

log_info "清理 Docker 构建缓存（24h 前）..."
docker builder prune -f --filter "until=24h" || true

echo ""
log_info "========================================="
log_info "清理完成"
log_info "========================================="
echo ""
log_info "清理后磁盘使用情况:"
df -h | grep -E "Filesystem|/$" || true
echo ""
log_info "清理后 Docker 磁盘使用:"
docker system df