#!/bin/bash

# Docker Swarm 管理脚本
set -euo pipefail

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查是否在 Swarm 模式
check_swarm_mode() {
    if docker info --format '{{.Swarm.LocalNodeState}}' | grep -q "active"; then
        log_success "Docker Swarm 模式已激活"
        return 0
    else
        log_error "Docker Swarm 模式未激活"
        return 1
    fi
}

# 初始化 Swarm
init_swarm() {
    log_info "初始化 Docker Swarm..."
    
    if docker swarm init; then
        log_success "Docker Swarm 初始化成功"
        log_info "要添加工作节点，请在其他主机上运行:"
        docker swarm join-token worker
    else
        log_error "Docker Swarm 初始化失败"
        return 1
    fi
}

# 部署堆栈
deploy_stack() {
    local stack_name="$1"
    local compose_file="$2"
    
    log_info "部署堆栈: $stack_name"
    
    if [[ ! -f "$compose_file" ]]; then
        log_error "Compose 文件不存在: $compose_file"
        return 1
    fi
    
    if docker stack deploy -c "$compose_file" "$stack_name"; then
        log_success "堆栈 $stack_name 部署成功"
    else
        log_error "堆栈 $stack_name 部署失败"
        return 1
    fi
}

# 列出堆栈
list_stacks() {
    log_info "当前部署的堆栈:"
    docker stack ls
}

# 查看堆栈服务
show_stack_services() {
    local stack_name="$1"
    
    log_info "堆栈 $stack_name 的服务:"
    docker stack services "$stack_name"
}

# 删除堆栈
remove_stack() {
    local stack_name="$1"
    
    log_info "删除堆栈: $stack_name"
    
    if docker stack rm "$stack_name"; then
        log_success "堆栈 $stack_name 删除成功"
    else
        log_error "堆栈 $stack_name 删除失败"
        return 1
    fi
}

# 显示节点信息
show_nodes() {
    log_info "Swarm 节点信息:"
    docker node ls
}

# 显示帮助信息
show_help() {
    echo "Docker Swarm 管理脚本"
    echo ""
    echo "用法: $0 [命令] [参数]"
    echo ""
    echo "命令:"
    echo "  init                          - 初始化 Docker Swarm"
    echo "  deploy <stack> <compose-file> - 部署堆栈"
    echo "  list                          - 列出所有堆栈"
    echo "  services <stack>              - 查看堆栈服务"
    echo "  remove <stack>                - 删除堆栈"
    echo "  nodes                         - 显示节点信息"
    echo "  check                         - 检查 Swarm 状态"
    echo "  help                          - 显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 init                                    # 初始化 Swarm"
    echo "  $0 deploy traefik stacks/traefik-swarm-stack.yml"
    echo "  $0 deploy demo stacks/demo-services-stack.yml"
    echo "  $0 list                                    # 列出堆栈"
    echo "  $0 services traefik                       # 查看 traefik 堆栈服务"
}

# 主函数
main() {
    case "${1:-help}" in
        "init")
            init_swarm
            ;;
        "deploy")
            if [[ $# -lt 3 ]]; then
                log_error "部署命令需要堆栈名称和 compose 文件"
                echo "用法: $0 deploy <stack-name> <compose-file>"
                exit 1
            fi
            check_swarm_mode || exit 1
            deploy_stack "$2" "$3"
            ;;
        "list")
            check_swarm_mode || exit 1
            list_stacks
            ;;
        "services")
            if [[ $# -lt 2 ]]; then
                log_error "需要指定堆栈名称"
                echo "用法: $0 services <stack-name>"
                exit 1
            fi
            check_swarm_mode || exit 1
            show_stack_services "$2"
            ;;
        "remove")
            if [[ $# -lt 2 ]]; then
                log_error "需要指定堆栈名称"
                echo "用法: $0 remove <stack-name>"
                exit 1
            fi
            check_swarm_mode || exit 1
            remove_stack "$2"
            ;;
        "nodes")
            check_swarm_mode || exit 1
            show_nodes
            ;;
        "check")
            check_swarm_mode
            ;;
        "help"|*)
            show_help
            ;;
    esac
}

# 运行主函数
main "$@"