#!/bin/bash
# Docker 缓存清理脚本
# 用法：
#   ./docker-cache-cleanup.sh                    # 默认清理（保留 3 小时内的缓存）
#   ./docker-cache-cleanup.sh --aggressive       # 激进清理（保留 1 小时内）
#   ./docker-cache-cleanup.sh --conservative     # 保守清理（保留 6 小时内）
#   ./docker-cache-cleanup.sh --full             # 完全清理（删除所有缓存）

set -euo pipefail

# ============================================================
# 配置
# ============================================================
REMOTE_HOST="${REMOTE_HOST:-txjp}"
CLEANUP_MODE="${1:---default}"
RETENTION_HOURS=3  # 默认保留 3 小时

# 生产环境服务器 IP 列表（用于判断是否在服务器本地运行）
# 新增服务器时，只需在此添加 IP 地址即可
PRODUCTION_SERVER_IPS=(
    "43.133.38.210"    # txjp (腾讯云北京)
    # "159.75.83.209"  # tgz (腾讯云广州) - 如需添加
    # "10.0.0.1"       # 其他服务器 - 如需添加
)

# 检测运行环境：
# 通过 curl ip.sb 获取当前公网 IP，与生产环境服务器 IP 列表匹配
IS_LOCAL=false

# 获取当前公网 IP
CURRENT_IP=$(curl -s --connect-timeout 5 ip.sb 2>/dev/null || echo "")

if [ -n "$CURRENT_IP" ]; then
    # 遍历生产环境服务器 IP 列表，检查是否匹配
    for SERVER_IP in "${PRODUCTION_SERVER_IPS[@]}"; do
        if [ "$CURRENT_IP" = "$SERVER_IP" ]; then
            IS_LOCAL=true
            break
        fi
    done
fi

# 如果是本地运行，不需要 SSH
if [ "$IS_LOCAL" = true ]; then
    REMOTE_HOST=""
fi

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

# ============================================================
# 日志函数
# ============================================================
log() { echo -e "${GREEN}[$(date '+%H:%M:%S')]${NC} $*"; }
warn() { echo -e "${YELLOW}[$(date '+%H:%M:%S')] WARNING:${NC} $*"; }
error() { echo -e "${RED}[$(date '+%H:%M:%S')] ERROR:${NC} $*" >&2; }

# ============================================================
# 清理模式设置
# ============================================================
case "$CLEANUP_MODE" in
    --aggressive)
        RETENTION_HOURS=1
        log "激进清理模式：保留 ${RETENTION_HOURS} 小时内的缓存"
        ;;
    --conservative)
        RETENTION_HOURS=6
        log "保守清理模式：保留 ${RETENTION_HOURS} 小时内的缓存"
        ;;
    --full)
        log "完全清理模式：删除所有缓存"
        RETENTION_HOURS=0
        ;;
    *)
        log "默认清理模式：保留 ${RETENTION_HOURS} 小时内的缓存"
        ;;
esac

# ============================================================
# 辅助函数
# ============================================================
# 执行命令（根据是否本地决定是否使用 SSH）
run_cmd() {
    if [ "$IS_LOCAL" = true ]; then
        eval "$@"
    else
        ssh "$REMOTE_HOST" "$@"
    fi
}

# 获取磁盘使用率
get_disk_usage() {
    run_cmd "df / --output=pcent | tail -1 | tr -d ' %'" 2>/dev/null || echo "0"
}

# 获取构建缓存大小
get_cache_size() {
    run_cmd "docker system df --format '{{.Size}}' 2>/dev/null | head -1" || echo "0"
}

# 获取镜像数量
get_image_count() {
    run_cmd "docker images -q 2>/dev/null | wc -l" || echo "0"
}

# ============================================================
# 清理函数
# ============================================================
# 清理 dangling 镜像（无标签或未被引用的中间层）
cleanup_dangling_images() {
    log "清理 dangling 镜像..."
    run_cmd "docker image prune -f 2>/dev/null" || warn "docker image prune 失败"
}

# 清理测试镜像（以 test- 开头的镜像）
cleanup_test_images() {
    log "清理测试镜像..."
    local test_images=$(run_cmd "docker images --format '{{.Repository}}:{{.Tag}}' | grep -E '^test-' 2>/dev/null" || true)
    if [ -n "$test_images" ]; then
        echo "$test_images" | while read -r img; do
            log "  删除: $img"
        done
        run_cmd "docker images --format '{{.Repository}}:{{.Tag}}' | grep -E '^test-' | xargs -r docker rmi -f 2>/dev/null" || true
    else
        log "  无测试镜像"
    fi
}

# 清理构建缓存
cleanup_build_cache() {
    if [ "$RETENTION_HOURS" -eq 0 ]; then
        log "清理所有构建缓存..."
        run_cmd "docker builder prune -af 2>/dev/null" || warn "docker builder prune 失败"
    else
        log "清理 ${RETENTION_HOURS} 小时前的构建缓存..."
        run_cmd "docker builder prune -f --filter 'until=${RETENTION_HOURS}h' 2>/dev/null" || warn "docker builder prune 失败"
    fi
}

# 根据磁盘使用率决定是否需要深度清理
cleanup_by_disk_usage() {
    local DISK_USAGE=$(get_disk_usage)
    log "当前磁盘使用率: ${DISK_USAGE}%"

    if [ "$DISK_USAGE" -gt 85 ]; then
        warn "磁盘使用率 ${DISK_USAGE}% > 85%，执行深度清理"
        # 清理所有构建缓存
        run_cmd "docker builder prune -af 2>/dev/null" || true
        # 清理容器日志
        log "截断容器日志..."
        run_cmd "truncate -s 0 /var/lib/docker/containers/\$(docker ps -q)/\*-json.log 2>/dev/null" || true
    elif [ "$DISK_USAGE" -gt 75 ]; then
        cleanup_build_cache
    else
        log "磁盘使用率正常，仅执行常规清理"
        cleanup_build_cache
    fi
}

# ============================================================
# 打印清理结果
# ============================================================
print_result() {
    local DISK_AFTER=$(get_disk_usage)
    local CACHE_AFTER=$(get_cache_size)
    local IMAGE_AFTER=$(get_image_count)
    echo ""
    log "========== 清理完成 =========="
    log "磁盘使用率: ${DISK_AFTER}%"
    log "构建缓存大小: ${CACHE_AFTER}"
    log "镜像数量: ${IMAGE_AFTER}"
}

# ============================================================
# 主函数
# ============================================================
main() {
    local DISK_BEFORE=$(get_disk_usage)
    local CACHE_BEFORE=$(get_cache_size)
    local IMAGE_BEFORE=$(get_image_count)

    echo ""
    log "========== Docker 缓存清理 =========="
    log "清理前 - 磁盘使用率: ${DISK_BEFORE}%, 缓存大小: ${CACHE_BEFORE}, 镜像数量: ${IMAGE_BEFORE}"
    echo ""

    # 执行清理
    cleanup_dangling_images
    cleanup_test_images
    cleanup_by_disk_usage

    # 打印结果
    print_result
}

# 执行主函数
main "$@"