FranklinWH-Toolbox/service/repositories/mqtt/fwh_mqtt/fwh_mqtt_client.go

package fwh_mqtt

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"fwh-toolbox-server/config"
	"fwh-toolbox-server/utils/mqttClient"
	"os"
	"strconv"
	"sync"
	"time"

	"go.uber.org/zap"
)

const (
	fwhClientIdPrefix  = "fwh_xlog_mqtt_client_"
	sleepTimeout       = 5 * time.Minute
	sendMessageTimeout = 1 * time.Minute
	subscriberBuffer   = 64
)

// transport 底层 MQTT 连接（utils/mqttClient），抽象出来便于测试
type transport interface {
	Connect() error
	IsConnected() bool
	Disconnect()
	SubscribeMessage(topic string, callback func(payload []byte)) error
	UnSubscribeMessage(topic string) error
	SendMessageOnce(topic string, msg []byte) error
	GetClientStatus() string
}

// Message 收到的一条消息
type Message struct {
	Topic      string
	Envelope   FwhMqttMsg // 外层；DataArea 为通用解码结果
	Raw        []byte     // 原始字节，业务层按需重新解析（如重复键）
	ReceivedAt time.Time
}

// Handler 订阅回调；每个订阅者在独立 goroutine 中顺序调用，可以阻塞，不影响其他订阅者
type Handler func(Message)

type subscriber struct {
	id      uint64
	handler Handler
	ch      chan Message
	done    chan struct{}
}

type topicEntry struct {
	subs  map[uint64]*subscriber
	ready chan struct{} // broker 订阅完成后关闭
	err   error
}

// opKind worker 串行执行的底层操作
type opKind int

const (
	opPublish opKind = iota
	opSubscribe
	opUnsubscribe
)

type op struct {
	kind    opKind
	topic   string
	payload []byte
	result  chan error
}

// FwhMqttClient 一个环境的 FWH MQTT 客户端。
// 只提供底层收发：Request（一问一答）、Publish（只请求）、Subscribe（持续监听），不包含任何业务判断。
type FwhMqttClient struct {
	env Env
	tr  transport
	ops chan op

	mu     sync.Mutex
	topics map[string]*topicEntry
	nextID uint64
}

func newFwhMqttClient(env Env, tr transport) *FwhMqttClient {
	c := &FwhMqttClient{
		env:    env,
		tr:     tr,
		ops:    make(chan op, 1000),
		topics: map[string]*topicEntry{},
	}
	go c.worker()
	return c
}

// Env FWH 服务器环境，每个环境使用独立的 broker、证书和客户端实例
type Env string

const (
	EnvProd Env = "prod" // 生产服
	EnvTest Env = "test" // 测试服
)

// ParseEnv 解析请求中的环境参数，空值默认为生产服
func ParseEnv(s string) (Env, error) {
	switch Env(s) {
	case "", EnvProd:
		return EnvProd, nil
	case EnvTest:
		return EnvTest, nil
	default:
		return "", fmt.Errorf("unknown FWH MQTT env: %s", s)
	}
}

var (
	fwhMqttClients   = map[Env]*FwhMqttClient{}
	fwhMqttClientsMu sync.RWMutex
)

// InitFwhMqttClient 初始化指定环境的 MQTT 客户端；失败时该环境不可用，不影响其他环境
func InitFwhMqttClient(env Env, cfg *config.FwhMqttConfig) error {
	if cfg.FwhMqttBroker == "" {
		zap.L().Warn("FWH MQTT broker not configured, skip init.", zap.String("env", string(env)))
		return fmt.Errorf("FWH MQTT broker not configured (env=%s)", env)
	}

	// 证书配置
	caCert, err := os.ReadFile(cfg.FwhMqttCaUrl)
	if err != nil {
		zap.L().Error("Read FWH MQTT CA failed.", zap.String("env", string(env)), zap.String("err", err.Error()))
		return err
	}

	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	cert, err := tls.LoadX509KeyPair(cfg.FwhMqttCertUrl, cfg.FwhMqttKeyUrl)
	if err != nil {
		zap.L().Error("Load FWH MQTT cert failed.", zap.String("env", string(env)), zap.String("err", err.Error()))
		return err
	}

	client, err := mqttClient.CreateMqttClient(mqttClient.MqttClientOpt{
		Broker:   cfg.FwhMqttBroker,
		ClientID: fwhClientIdPrefix + strconv.FormatInt(time.Now().Unix(), 10),
		TLS: &tls.Config{
			Certificates: []tls.Certificate{cert},
			RootCAs:      caCertPool,
		},
	})
	if err != nil {
		zap.L().Error("Init FWH MQTT client failed.", zap.String("env", string(env)))
		return err
	}

	return registerClient(env, client)
}

func registerClient(env Env, tr transport) error {
	fwhMqttClientsMu.Lock()
	defer fwhMqttClientsMu.Unlock()
	if _, exists := fwhMqttClients[env]; exists {
		return fmt.Errorf("FWH MQTT client already initialized (env=%s)", env)
	}
	fwhMqttClients[env] = newFwhMqttClient(env, tr)
	zap.L().Info("Init FWH MQTT client success.", zap.String("env", string(env)))
	return nil
}

// GetClient 获取指定环境的 MQTT 客户端
func GetClient(env Env) (*FwhMqttClient, error) {
	fwhMqttClientsMu.RLock()
	defer fwhMqttClientsMu.RUnlock()

	c, ok := fwhMqttClients[env]
	if !ok {
		return nil, fmt.Errorf("Fwh Mqtt Client Uninitialized (env=%s).", env)
	}
	return c, nil
}

/************** worker：独占底层连接，串行执行 **************/

func (c *FwhMqttClient) worker() {
	idleTimer := time.NewTimer(sleepTimeout)
	defer idleTimer.Stop()

	for {
		select {
		case o := <-c.ops:
			if !idleTimer.Stop() {
				select {
				case <-idleTimer.C:
				default:
				}
			}
			idleTimer.Reset(sleepTimeout)
			o.result <- c.execute(o)

		case <-idleTimer.C:
			// 有活动订阅时保持连接，否则空闲断开以节省连接
			if c.tr.IsConnected() && !c.hasTopics() {
				zap.L().Info("FWH MQTT no task received within timeout, disconnect.", zap.String("env", string(c.env)))
				c.tr.Disconnect()
			}
			idleTimer.Reset(sleepTimeout)
		}
	}
}

func (c *FwhMqttClient) execute(o op) error {
	if err := c.tr.Connect(); err != nil {
		zap.L().Error("FWH MQTT connect failed.",
			zap.String("env", string(c.env)),
			zap.String("status", c.tr.GetClientStatus()),
			zap.String("err", err.Error()),
		)
		return fmt.Errorf("FWH MQTT connect failed: %w", err)
	}

	switch o.kind {
	case opPublish:
		return c.tr.SendMessageOnce(o.topic, o.payload)
	case opSubscribe:
		topic := o.topic
		return c.tr.SubscribeMessage(topic, func(payload []byte) { c.dispatch(topic, payload) })
	case opUnsubscribe:
		return c.tr.UnSubscribeMessage(o.topic)
	}
	return fmt.Errorf("unknown op %d", o.kind)
}

// enqueue 提交操作，返回结果 channel
func (c *FwhMqttClient) enqueue(kind opKind, topic string, payload []byte) chan error {
	o := op{kind: kind, topic: topic, payload: payload, result: make(chan error, 1)}
	c.ops <- o
	return o.result
}

func waitResult(ctx context.Context, ch chan error) error {
	select {
	case err := <-ch:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (c *FwhMqttClient) hasTopics() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.topics) > 0
}

/************** 对外方法 **************/

// Publish 只请求：发送后立即返回
func (c *FwhMqttClient) Publish(ctx context.Context, topic string, msg FwhMqttMsg) error {
	b, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(ctx, sendMessageTimeout)
	defer cancel()
	if err := waitResult(ctx, c.enqueue(opPublish, topic, b)); err != nil {
		zap.L().Error("FWH MQTT publish failed.", zap.String("env", string(c.env)), zap.String("topic", topic), zap.String("err", err.Error()))
		return err
	}
	return nil
}

// Subscribe 持续监听：同一 topic 可有多个订阅者，各自收到每条消息；返回的 cancel 只移除自己。
// 首个订阅者加入时向 broker 订阅，最后一个退出时取消订阅；有订阅时连接不会空闲断开。
func (c *FwhMqttClient) Subscribe(ctx context.Context, topic string, handler Handler) (cancel func(), err error) {
	c.mu.Lock()
	c.nextID++
	sub := &subscriber{
		id:      c.nextID,
		handler: handler,
		ch:      make(chan Message, subscriberBuffer),
		done:    make(chan struct{}),
	}
	entry, exists := c.topics[topic]
	var result chan error
	if !exists {
		entry = &topicEntry{subs: map[uint64]*subscriber{}, ready: make(chan struct{})}
		c.topics[topic] = entry
		result = c.enqueue(opSubscribe, topic, nil) // 持锁入队，保证与取消订阅的先后顺序
	}
	entry.subs[sub.id] = sub
	c.mu.Unlock()

	if !exists {
		err := waitResult(ctx, result)
		c.mu.Lock()
		entry.err = err
		close(entry.ready)
		if err != nil && c.topics[topic] == entry {
			delete(c.topics, topic)
			// 等待超时时 broker 订阅可能稍后成功，补发取消订阅避免残留
			c.enqueue(opUnsubscribe, topic, nil)
		}
		c.mu.Unlock()
	} else {
		select {
		case <-entry.ready:
		case <-ctx.Done():
			c.removeSubscriber(topic, sub)
			return nil, ctx.Err()
		}
	}
	if entry.err != nil {
		zap.L().Error("FWH MQTT subscribe failed.", zap.String("env", string(c.env)), zap.String("topic", topic), zap.String("err", entry.err.Error()))
		return nil, fmt.Errorf("subscribe failed: %w", entry.err)
	}

	go func() {
		for {
			select {
			case m := <-sub.ch:
				sub.handler(m)
			case <-sub.done:
				return
			}
		}
	}()

	var once sync.Once
	return func() { once.Do(func() { c.removeSubscriber(topic, sub) }) }, nil
}

func (c *FwhMqttClient) removeSubscriber(topic string, sub *subscriber) {
	c.mu.Lock()
	defer c.mu.Unlock()
	close(sub.done)
	entry, ok := c.topics[topic]
	if !ok {
		return
	}
	delete(entry.subs, sub.id)
	if len(entry.subs) == 0 {
		delete(c.topics, topic)
		result := c.enqueue(opUnsubscribe, topic, nil)
		go func() {
			if err := <-result; err != nil {
				zap.L().Warn("FWH MQTT unsubscribe failed.", zap.String("env", string(c.env)), zap.String("topic", topic), zap.String("err", err.Error()))
			}
		}()
	}
}

// dispatch 在 paho 回调线程中执行：解析一次外层，非阻塞分发给各订阅者
func (c *FwhMqttClient) dispatch(topic string, payload []byte) {
	m := Message{Topic: topic, Raw: payload, ReceivedAt: time.Now()}
	if err := json.Unmarshal(payload, &m.Envelope); err != nil {
		zap.L().Warn("FWH MQTT message unmarshal failed.", zap.String("topic", topic), zap.String("payload", string(payload)), zap.String("err", err.Error()))
		return
	}

	c.mu.Lock()
	entry := c.topics[topic]
	var subs []*subscriber
	if entry != nil {
		for _, s := range entry.subs {
			subs = append(subs, s)
		}
	}
	c.mu.Unlock()

	for _, s := range subs {
		select {
		case <-s.done:
		case s.ch <- m:
		default:
			zap.L().Warn("FWH MQTT subscriber buffer full, dropping message.", zap.String("topic", topic), zap.Uint("cmdType", m.Envelope.CmdType))
		}
	}
}

// RequestOpt 一问一答参数
type RequestOpt struct {
	Topic    string             // 发送 topic
	Msg      FwhMqttMsg         // 发送内容
	ResTopic string             // 应答 topic
	Match    func(Message) bool // 应答匹配条件，由调用方定义；为 nil 时取第一条消息
	Timeout  time.Duration      // 等待应答超时，默认 sendMessageTimeout
}

// Request 一问一答：先订阅应答 topic，再发送，返回第一条满足 Match 的消息
func (c *FwhMqttClient) Request(ctx context.Context, opt RequestOpt) (*Message, error) {
	timeout := opt.Timeout
	if timeout <= 0 {
		timeout = sendMessageTimeout
	}
	ctx, cancelCtx := context.WithTimeout(ctx, timeout)
	defer cancelCtx()

	resCh := make(chan Message, 1)
	cancelSub, err := c.Subscribe(ctx, opt.ResTopic, func(m Message) {
		if opt.Match != nil && !opt.Match(m) {
			return
		}
		select {
		case resCh <- m:
		default:
		}
	})
	if err != nil {
		return nil, err
	}
	defer cancelSub()

	if err := c.Publish(ctx, opt.Topic, opt.Msg); err != nil {
		return nil, err
	}

	select {
	case m := <-resCh:
		return &m, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("FWH MQTT received message timeout")
	}
}