common/message/rabbitmq_msg.go

package message

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/gogf/gf/v2/frame/g"
	amqp "github.com/rabbitmq/amqp091-go"
)

// rabbitMQMessageClient RabbitMQ 实现
type rabbitMQMessageClient struct {
	clientType messageClientType
}

// StreamGroup 创建消费组（支持单个或批量）
func (q *rabbitMQMessageClient) streamGroup(ctx context.Context, configs ...interface{}) error {
	if len(configs) == 0 {
		return fmt.Errorf("配置不能为空")
	}
	for _, config := range configs {
		cfg, ok := config.(*RabbitMQConfig)
		if !ok {
			return fmt.Errorf("无效的 RabbitMQ 配置类型")
		}
		if err := q.setupQueue(ctx, channel, cfg, cfg.DelayMessage); err != nil {
			return err
		}
	}
	return nil
}

// Publish 发布消息（支持单个或批量）
func (q *rabbitMQMessageClient) publish(ctx context.Context, config interface{}, data interface{}) error {
	cfg, ok := config.(*RabbitMQConfig)
	if !ok {
		return fmt.Errorf("无效的 RabbitMQ 配置类型")
	}
	if err := q.publishMessage(ctx, cfg, "work", data, 0); err != nil {
		g.Log().Errorf(ctx, "❌ RabbitMQ 发布消息失败: err=%v", err)
		return err
	}
	return nil
}

// PublishDelayed 发布延迟消息
func (q *rabbitMQMessageClient) publishDelayed(ctx context.Context, config interface{}, data interface{}, delaySeconds int) error {
	cfg, ok := config.(*RabbitMQConfig)
	if !ok {
		return fmt.Errorf("无效的 RabbitMQ 配置类型")
	}
	if err := q.publishMessage(ctx, cfg, "delayed", data, delaySeconds); err != nil {
		g.Log().Errorf(ctx, "❌ RabbitMQ 发布延迟消息失败: err=%v", err)
		return err
	}
	return nil
}

func (q *rabbitMQMessageClient) publishMessage(ctx context.Context, cfg *RabbitMQConfig, mode string, data interface{}, delaySeconds int) error {
	body, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("序列化数据失败: %w", err)
	}
	deliveryMode := amqp.Transient
	if cfg.Durable {
		deliveryMode = amqp.Persistent
	}
	publishing := amqp.Publishing{
		ContentType:  "application/json",
		Body:         body,
		DeliveryMode: deliveryMode,
		Timestamp:    time.Now(),
	}
	if delaySeconds > 0 {
		publishing.Headers = amqp.Table{
			"x-delay": delaySeconds * 1000, // 延时时间（毫秒）
		}
	}
	exchange, routingKey := q.parseExchangeAndRoutingKey(ctx, mode, cfg)
	err = channel.PublishWithContext(
		ctx,
		exchange,
		routingKey,
		false, false,
		publishing,
	)
	return err
}

func (q *rabbitMQMessageClient) parseExchangeAndRoutingKey(_ context.Context, mode string, cfg *RabbitMQConfig) (exchange, routingKey string) {
	switch mode {
	case "work", "":
		exchange = ""         // 默认交换机
		routingKey = cfg.Name // 队列名
	case "event", "topic":
		exchange = cfg.Exchange
		routingKey = cfg.Topic
	case "broadcast":
		exchange = cfg.Exchange
		routingKey = "" // fanout忽略路由键
	case "delayed":
		exchange = cfg.Exchange + ".delayed"
		routingKey = cfg.Topic
	default:
		exchange = ""
		routingKey = cfg.Name
	}
	return exchange, routingKey
}

// setupQueue 统一的队列设置方法（声明 Exchange、队列、绑定、延迟 Exchange）
func (q *rabbitMQMessageClient) setupQueue(ctx context.Context, ch *amqp.Channel, cfg *RabbitMQConfig, delayMessage bool) error {
	exchange, routingKey := q.parseExchangeAndRoutingKey(ctx, cfg.Mode, cfg)

	// 声明 Exchange
	if err := ch.ExchangeDeclare(exchange, "topic", cfg.Durable, false, false, false, nil); err != nil {
		return fmt.Errorf("声明 Exchange 失败: %w", err)
	}

	// 声明队列
	if _, err := ch.QueueDeclare(cfg.Queue, cfg.Durable, false, false, false, nil); err != nil {
		return fmt.Errorf("声明队列失败: %w", err)
	}

	// 绑定队列
	if err := ch.QueueBind(cfg.Queue, routingKey, exchange, false, nil); err != nil {
		return fmt.Errorf("绑定队列失败: %w", err)
	}

	// 声明延迟 Exchange（如果需要）
	if delayMessage {
		if err := ch.ExchangeDeclare(exchange, "x-delayed-message", true, false, false, false, amqp.Table{"x-delayed-type": "direct"}); err != nil {
			return fmt.Errorf("声明延迟 Exchange 失败: %w", err)
		}
		if err := ch.QueueBind(cfg.Name, routingKey, exchange, false, nil); err != nil {
			return fmt.Errorf("绑定延迟队列失败: %w", err)
		}
	}

	return nil
}

// Subscribe 订阅消息（支持单个或批量）
func (q *rabbitMQMessageClient) subscribe(ctx context.Context, configs ...interface{}) error {
	if len(configs) == 0 {
		return fmt.Errorf("配置不能为空")
	}

	for _, config := range configs {
		cfg, ok := config.(*RabbitMQConfig)
		if !ok {
			return fmt.Errorf("无效的 RabbitMQ 配置类型")
		}
		handler := cfg.HandleFunc
		if handler == nil {
			return fmt.Errorf("必须提供处理函数")
		}
		if err := q.createSubscribe(ctx, cfg, handler); err != nil {
			return err
		}
	}
	return nil
}

// subscribe 内部单个订阅消息
func (q *rabbitMQMessageClient) createSubscribe(ctx context.Context, cfg *RabbitMQConfig, handler func(ctx context.Context, message map[string]interface{}) error) error {
	g.Log().Infof(ctx, "🔔 RabbitMQ 开始订阅: exchange=%s, queue=%s", cfg.Exchange, cfg.Queue)

	// 设置 Qos (预取数量)，控制每次推送的消息数量
	// prefetchCount: 未 ACK 消息的最大数量
	// prefetchSize: 未 ACK 消息的总大小（0 表示不限制）
	// global: false 表示仅应用于当前消费者
	prefetchCount := cfg.PrefetchCount
	if prefetchCount <= 0 {
		prefetchCount = 10 // 默认值为 10
	}
	if err := channel.Qos(prefetchCount, 0, false); err != nil {
		return fmt.Errorf("设置 Qos 失败: %w", err)
	}
	g.Log().Infof(ctx, "📊 设置 Prefetch Count: %d", prefetchCount)

	msg, err := channel.Consume(
		cfg.Queue,   // queue
		cfg.Queue,   // consumer
		cfg.AutoAck, // auto-ack (根据配置决定)
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // args
	)
	if err != nil {
		return fmt.Errorf("注册消费者失败: %w", err)
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				g.Log().Errorf(ctx, "❌ RabbitMQ 消费者 panic: %v", r)
			}
		}()

		// 并发控制信号量
		semaphore := make(chan struct{}, 10) // 限制最大并发数为 10

		for {
			select {
			case <-ctx.Done():
				g.Log().Infof(ctx, "🔕 RabbitMQ 消费者停止: queue=%s", cfg.Queue)
				return
			case msg, ok := <-msg:
				if !ok {
					g.Log().Warningf(ctx, "⚠️ RabbitMQ 消息通道关闭")
					return
				}

				// 获取并发控制槽位
				semaphore <- struct{}{}

				go func(m amqp.Delivery) {
					defer func() {
						<-semaphore // 释放槽位
						if r := recover(); r != nil {
							g.Log().Errorf(ctx, "❌ 消息处理 panic: %v", r)
						}
					}()

					if err := q.handleMessageWithRetry(ctx, m, handler, cfg.MaxRetry); err != nil {
						g.Log().Errorf(ctx, "❌ 消息处理失败(重试次数耗尽): %v", err)

						// 仅在手动 ACK 模式下拒绝消息
						if !cfg.AutoAck {
							// 拒绝消息不再重新入队（避免死循环）
							m.Nack(false, false)
						}
						return
					}

					// 仅在手动 ACK 模式下确认消息
					if cfg.AutoAck {
						if err := m.Ack(false); err != nil {
							g.Log().Errorf(ctx, "❌ ACK 消息失败: %v", err)
						}
					}
				}(msg)
			}
		}
	}()

	return nil
}

// handleMessageWithRetry 处理消息（支持重试）
func (q *rabbitMQMessageClient) handleMessageWithRetry(ctx context.Context, msg amqp.Delivery, handler func(ctx context.Context, message map[string]interface{}) error, maxRetry int) error {
	var data map[string]interface{}

	if err := json.Unmarshal(msg.Body, &data); err != nil {
		// 如果不是 JSON，直接使用原始内容
		data = map[string]interface{}{
			"data": string(msg.Body),
		}
	}

	// 重试逻辑
	for attempt := 0; attempt <= maxRetry; attempt++ {
		if attempt > 0 {
			g.Log().Infof(ctx, "🔄 消息处理重试 (第%d次)", attempt)
			// 指数退避
			time.Sleep(time.Duration(attempt) * time.Second)
		}

		err := handler(ctx, data)
		if err == nil {
			return nil // 成功
		}

		g.Log().Warningf(ctx, "⚠️ 消息处理失败 (第%d次): %v", attempt+1, err)

		if attempt == maxRetry {
			return fmt.Errorf("达到最大重试次数 %d: %w", maxRetry, err)
		}
	}

	return nil
}