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

	topicPrefixS2C = "fwh/s2c/"
	topicPrefixC2S = "fwh/c2s/"
)

type FwhMqttResult struct {
	Err  error
	Data *FwhMqttMsg
}

type FwhMqttMsg struct {
	CmdType   uint        `json:"cmdType"`   // 业务类型标识
	Crc       string      `json:"crc"`       // dataArea 数据字符串的CRC
	DataArea  interface{} `json:"dataArea"`  // 数据域
	EquipNo   string      `json:"equipNo"`   // 设备编号 SN
	Len       int         `json:"len"`       // dataArea 数据长度
	Snno      uint        `json:"snno"`      // 序号
	TimeStamp int64       `json:"timeStamp"` // 时间戳
	Type      uint        `json:"type"`      // 设备类型 0：设备网关IBG
}

type MqttMsg struct {
	Topic      string // 发送 Topic
	Msg        FwhMqttMsg
	ResTopic   *string // 响应 Topic
	ResCmdType *uint
	ResultCh   chan FwhMqttResult
	Ctx        context.Context
}

type ResultChan struct {
	ch  chan FwhMqttMsg
	ctx context.Context
}

type FwhMqttClient struct {
	Client *mqttClient.MqttClient
	TaskCh chan MqttMsg
}

// 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
	}

	c := &FwhMqttClient{
		Client: client,
		TaskCh: make(chan MqttMsg, 1000),
	}

	fwhMqttClientsMu.Lock()
	if _, exists := fwhMqttClients[env]; exists {
		fwhMqttClientsMu.Unlock()
		return fmt.Errorf("FWH MQTT client already initialized (env=%s)", env)
	}
	fwhMqttClients[env] = c
	fwhMqttClientsMu.Unlock()

	zap.L().Info("Init FWH MQTT client success.", zap.String("env", string(env)))

	go c.fwhMqttWorker()

	return nil
}

// getFwhMqttClient 获取指定环境的 MQTT 客户端
func getFwhMqttClient(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
}

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

	for {
		select {
		case task, ok := <-c.TaskCh:
			if !ok {
				zap.L().Error("Channel closed, worker exiting, recreate.")
				c.TaskCh = make(chan MqttMsg, 1000)
				continue
			}

			// 重置空闲计时器
			if !idleTimer.Stop() {
				select {
				case <-idleTimer.C:
				default:
				}
			}
			idleTimer.Reset(sleepTimeout)

			// 处理任务
			c.handleTask(task)

		case <-idleTimer.C:
			if c.Client.IsConnected() {
				zap.L().Info("FWH MQTT no task received within timeout, disconnect.")
				c.Client.Disconnect()
			}
			idleTimer.Reset(sleepTimeout)
		}
	}
}

func (c *FwhMqttClient) handleTask(task MqttMsg) {
	// 1. 建立连接
	if err := c.Client.Connect(); err != nil {
		zap.L().Error("FWH MQTT connect failed.",
			zap.String("status", c.Client.GetClientStatus()),
			zap.String("err", err.Error()),
		)
		task.ResultCh <- FwhMqttResult{Err: fmt.Errorf("FWH MQTT connect failed: %w", err)}
		return
	}

	// 2. 如果需要响应，订阅响应 topic
	var resultCh chan FwhMqttResult
	if task.ResTopic != nil {
		resultCh = make(chan FwhMqttResult, 1)
		if err := c.Client.SubscribeMessage(*task.ResTopic, func(payload []byte) {
			var fwhMqttMsg FwhMqttMsg
			if err := json.Unmarshal(payload, &fwhMqttMsg); err != nil {
				zap.L().Error("Result Unmarshal failed.(FwhMqttMsg)",
					zap.String("topic", task.Topic),
					zap.String("payload", string(payload)),
					zap.String("err", err.Error()),
				)
				return
			}

			// 过滤 CmdType
			if task.ResCmdType != nil && fwhMqttMsg.CmdType != *task.ResCmdType {
				return
			}

			// 非阻塞写入，防止回调卡死
			select {
			case resultCh <- FwhMqttResult{Err: nil, Data: &fwhMqttMsg}:
			default:
				zap.L().Warn("[MQTT] resultCh full, dropping duplicate response")
			}
		}); err != nil {
			zap.L().Error("FWH MQTT subscribe failed.", zap.String("err", err.Error()), zap.String("topic", *task.ResTopic))
			task.ResultCh <- FwhMqttResult{Err: fmt.Errorf("subscribe failed: %w", err)}
			return
		}

		// 确保退出时取消订阅
		defer func() {
			if err := c.Client.UnSubscribeMessage(*task.ResTopic); err != nil {
				zap.L().Error("FWH MQTT unsubscribe failed.", zap.String("err", err.Error()), zap.String("topic", *task.ResTopic))
			}
		}()
	}

	// 3. 发送消息
	msgByte, err := json.Marshal(task.Msg)
	if err != nil {
		zap.L().Error("FWH MQTT message marshal failed.", zap.String("err", err.Error()))
		task.ResultCh <- FwhMqttResult{Err: err}
		return
	}

	if err := c.Client.SendMessageOnce(task.Topic, msgByte); err != nil {
		zap.L().Error("FWH MQTT send message failed.", zap.String("err", err.Error()))
		task.ResultCh <- FwhMqttResult{Err: err}
		return
	}

	// 4. 等待响应 或 直接返回
	if resultCh == nil {
		// 只发送，不等响应
		task.ResultCh <- FwhMqttResult{Err: nil, Data: nil}
		return
	}

	// 一问一答：等待响应或超时
	select {
	case res := <-resultCh:
		task.ResultCh <- res
	case <-task.Ctx.Done():
		task.ResultCh <- FwhMqttResult{Err: fmt.Errorf("FWH MQTT received message timeout")}
	}
}

func (c *FwhMqttClient) sendMessage(ctx context.Context, topicName string, msg FwhMqttMsg, resTopicName *string, resCmdType *uint) (*FwhMqttMsg, error) {
	if c.Client == nil {
		return nil, fmt.Errorf("Fwh Mqtt Client Uninitialized.")
	}

	ctxSend, cancel := context.WithTimeout(ctx, sendMessageTimeout)
	defer cancel()

	topic := topicPrefixS2C + topicName
	var resTopic *string
	if resTopicName != nil {
		resTopicStr := topicPrefixC2S + *resTopicName
		resTopic = &resTopicStr
	}

	taskMsg := MqttMsg{
		Topic:      topic,
		Msg:        msg,
		ResTopic:   resTopic,
		ResCmdType: resCmdType,
		ResultCh:   make(chan FwhMqttResult, 1),
		Ctx:        ctxSend,
	}
	c.TaskCh <- taskMsg

	for {
		select {
		case result, ok := <-taskMsg.ResultCh:
			if !ok {
				zap.L().Error("Result channel closed, worker exiting, recreate.", zap.String("topic", topic), zap.String("msg", fmt.Sprintf("%+v", msg)))
				return nil, fmt.Errorf("Result channel closed, worker exiting, recreate.")
			}

			if result.Err != nil {
				return nil, result.Err
			}

			return result.Data, nil
		case <-ctxSend.Done():
			zap.L().Info("Result received timeout.", zap.String("topic", topic), zap.String("msg", fmt.Sprintf("%+v", msg)))
			return nil, fmt.Errorf("Result received timeout.")
		}
	}
}

// 仅发送不等待响应结果
func (c *FwhMqttClient) sendMessageOnce(ctx context.Context, topicName string, msg FwhMqttMsg) error {
	_, err := c.sendMessage(ctx, topicName, msg, nil, nil)
	return err
}

// 发送消息后，等待响应结果
func (c *FwhMqttClient) sendMessageWait(ctx context.Context, topicName string, msg FwhMqttMsg, resTopicName string, resCmdType uint) (*FwhMqttMsg, error) {
	res, err := c.sendMessage(ctx, topicName, msg, &resTopicName, &resCmdType)
	return res, err
}