common/message/nats_msg.go

package message

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

	"github.com/gogf/gf/v2/frame/g"
	"github.com/nats-io/nats.go/jetstream"
)

// natsMessageClient NATS 实现
type natsMessageClient struct {
	clientType messageClientType
}

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

// createStreamGroup 内部单个创建消费组
func (q *natsMessageClient) createStreamGroup(ctx context.Context, cfg *NATSConfig) error {
	// Stream 不存在，创建新的
	storage := jetstream.FileStorage
	if !cfg.Durable {
		storage = jetstream.MemoryStorage
	}
	if g.IsEmpty(cfg.Replicas) {
		cfg.Replicas = 1
	}
	// 构建流配置
	jsConfig := jetstream.StreamConfig{
		Name:              cfg.Stream,
		Subjects:          []string{fmt.Sprintf("%s.>", cfg.Stream)},
		Replicas:          cfg.Replicas,
		NoAck:             cfg.AutoAck,
		AllowMsgSchedules: cfg.DelayMessage, // 延迟消息核心开关
		Storage:           storage,
		Discard:           jetstream.DiscardOld, // 达到上限删除旧消息
	}
	// 检查流是否已存在
	stream, err := js.Stream(ctx, cfg.Stream)
	if err == nil {
		// 流已存在，更新配置
		_, err = js.UpdateStream(ctx, jsConfig)
		if err != nil {
			return fmt.Errorf("更新任务流失败: %w", err)
		}
		g.Log().Infof(ctx, "任务流已更新: %s", stream.CachedInfo().Config.Name)
		return nil
	}
	// 创建新流
	stream, err = js.CreateStream(ctx, jsConfig)
	if err != nil {
		return fmt.Errorf("创建任务流失败: %w", err)
	}

	g.Log().Infof(ctx, "✅ NATS 队列初始化成功: stream=%s, consumer=%s", cfg.Stream, cfg.Consumer)
	return nil
}

// Publish 发布消息（支持单个或批量）
func (q *natsMessageClient) publish(ctx context.Context, config interface{}, data interface{}) error {
	cfg, ok := config.(*NATSConfig)
	if !ok {
		return fmt.Errorf("无效的 NATS 配置类型")
	}
	err := q.createStreamGroup(ctx, cfg)
	if err != nil {
		return err
	}
	payload, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("序列化数据失败: %w", err)
	}

	// 发布消息到 JetStream
	subject := fmt.Sprintf("%s.>", cfg.Stream)
	_, err = js.Publish(ctx, subject, payload)
	if err != nil {
		g.Log().Errorf(ctx, "❌ NATS 发布消息失败: topic=%s, err=%v", cfg.Stream, err)
		return err
	}

	g.Log().Infof(ctx, "✅ NATS 发布消息成功: topic=%s", cfg.Stream)
	return nil
}

// PublishDelayed 发布延迟消息（支持单个或批量）
func (q *natsMessageClient) publishDelayed(ctx context.Context, config interface{}, data interface{}, delay int) error {

	cfg, ok := config.(*NATSConfig)
	if !ok {
		return fmt.Errorf("无效的 NATS 配置类型")
	}

	payload, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("序列化数据失败: %w", err)
	}

	// 使用 goroutine 实现简单的延迟发布
	go func() {
		time.Sleep(time.Duration(delay))
		subject := fmt.Sprintf("%s.>", cfg.Stream)
		if err := q.publishInternal(ctx, subject, payload); err != nil {
			g.Log().Errorf(ctx, "❌ NATS 延迟消息发布失败: topic=%s, delay=%v, err=%v", cfg.Stream, delay, err)
		}
	}()

	g.Log().Infof(ctx, "✅ NATS 延迟消息已提交: topic=%s, delay=%v", cfg.Stream, delay)

	return nil
}

// publishInternal 内部发布消息
func (q *natsMessageClient) publishInternal(ctx context.Context, subject string, payload []byte) error {
	_, err := js.Publish(ctx, subject, payload)
	return err
}

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

	for _, config := range configs {
		cfg, ok := config.(*NATSConfig)
		if !ok {
			return fmt.Errorf("无效的 NATS 配置类型")
		}
		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 *natsMessageClient) createSubscribe(ctx context.Context, cfg *NATSConfig, handler func(ctx context.Context, message map[string]interface{}) error) error {
	g.Log().Infof(ctx, "🔔 NATS 开始订阅: stream=%s, consumer=%s", cfg.Stream, cfg.Consumer)
	// Stream 不存在，创建新的
	ackPolicy := jetstream.AckExplicitPolicy
	if cfg.AutoAck {
		ackPolicy = jetstream.AckNonePolicy
	}
	jsConfig := jetstream.ConsumerConfig{
		Name:          cfg.Consumer,
		Durable:       cfg.Consumer,
		AckPolicy:     ackPolicy,
		MaxDeliver:    3,
		MaxAckPending: cfg.PrefetchCount,
	}
	// 创建新消费者
	consumer, err := js.CreateOrUpdateConsumer(ctx, cfg.Stream, jsConfig)
	if err != nil {
		return fmt.Errorf("创建消费者失败: %w", err)
	}
	// 创建消息处理函数
	msgHandler := func(msg jetstream.Msg) {
		// 解析消息
		var data map[string]any
		if err := json.Unmarshal(msg.Data(), &data); err != nil {
			g.Log().Errorf(ctx, "解析消息失败: %v", err)
			msg.Nak()
			return
		}
		// 处理业务逻辑
		if err := handler(ctx, data); err != nil {
			g.Log().Errorf(ctx, "处理消息失败: %v", err)
			msg.Nak()
			return
		}
		g.Log().Infof(ctx, "处理消息成功")
		if !cfg.AutoAck {
			msg.Ack()
		}
	}

	// 开始消费
	_, err = consumer.Consume(msgHandler)
	if err != nil {
		return fmt.Errorf("开始消费失败: %w", err)
	}

	g.Log().Infof(ctx, "✅ 开始消费消息: %s/%s", cfg.Stream, cfg.Consumer)

	return nil
}