exchange_go/services/bitgetservice/bitget_market_client.go

package bitgetservice

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"go-admin/services/proxy"
	"strings"
	"sync"
	"time"

	"github.com/bytedance/sonic"
	log "github.com/go-admin-team/go-admin-core/logger"
	"github.com/gorilla/websocket"
)

// BitgetMarketClient Bitget公开行情订阅客户端
// 基于官方SDK设计，专门用于公开市场数据订阅
type BitgetMarketClient struct {
	config         *BitgetMarketConfig
	conn           *websocket.Conn
	subscriptions  map[SubscribeReq]OnReceive
	allSubscribes  *Set
	listener       OnReceive
	errorListener  OnReceive
	sendMutex      sync.Mutex
	reconnectMutex sync.Mutex
	ctx            context.Context
	cancel         context.CancelFunc
	connected      bool
	lastPongTime   time.Time
	proxyType      string
	proxyAddress   string
	reconnectCount int
}

// NewBitgetMarketClient 创建新的Bitget市场数据客户端
func NewBitgetMarketClient(config *BitgetMarketConfig) *BitgetMarketClient {
	if config == nil {
		config = DefaultMarketConfig()
	}

	ctx, cancel := context.WithCancel(context.Background())

	return &BitgetMarketClient{
		config:        config,
		subscriptions: make(map[SubscribeReq]OnReceive),
		allSubscribes: NewSet(),
		ctx:           ctx,
		cancel:        cancel,
		connected:     false,
		lastPongTime:  time.Now(),
	}
}

// SetProxy 设置代理
func (c *BitgetMarketClient) SetProxy(proxyType, proxyAddress string) {
	c.proxyType = proxyType
	c.proxyAddress = proxyAddress
}

// SetListeners 设置消息监听器
func (c *BitgetMarketClient) SetListeners(msgListener, errorListener OnReceive) {
	c.listener = msgListener
	c.errorListener = errorListener
}

// Connect 连接WebSocket
func (c *BitgetMarketClient) Connect() error {
	c.reconnectMutex.Lock()
	defer c.reconnectMutex.Unlock()

	if c.connected {
		return nil
	}

	// 创建WebSocket连接
	dialer := websocket.DefaultDialer

	// 设置代理
	if c.proxyAddress != "" {
		var err error
		dialer, err = proxy.GetDialer(c.proxyType, c.proxyAddress)
		if err != nil {
			return fmt.Errorf("设置代理失败: %v", err)
		}
	}

	dialer.HandshakeTimeout = 10 * time.Second
	log.Infof("正在连接Bitget市场数据WebSocket: %s", c.config.WsUrl)
	conn, _, err := dialer.Dial(c.config.WsUrl, nil)
	if err != nil {
		return fmt.Errorf("webSocket连接失败: %v", err)
	}

	c.conn = conn
	c.connected = true
	c.lastPongTime = time.Now()
	c.reconnectCount = 0

	log.Info("Bitget市场数据WebSocket连接成功")

	// 启动消息处理协程
	go c.readLoop()
	go c.pingLoop()
	go c.monitorConnection()

	return nil
}

// Subscribe 订阅市场数据
func (c *BitgetMarketClient) Subscribe(subscribeReqs []SubscribeReq, listener OnReceive) error {
	if !c.connected {
		if err := c.Connect(); err != nil {
			return err
		}
	}

	var args []interface{}
	for _, req := range subscribeReqs {
		// 标准化订阅请求
		normalizedReq := c.normalizeSubscribeReq(req)
		args = append(args, normalizedReq)

		// 保存订阅信息
		c.subscriptions[normalizedReq] = listener
		c.allSubscribes.Add(normalizedReq)
	}

	// 发送订阅消息
	wsReq := WsBaseReq{
		Op:   "subscribe",
		Args: args,
	}

	return c.sendMessage(wsReq)
}

// Unsubscribe 取消订阅
func (c *BitgetMarketClient) Unsubscribe(subscribeReqs []SubscribeReq) error {
	if !c.connected {
		return errors.New("WebSocket未连接")
	}

	var args []interface{}
	for _, req := range subscribeReqs {
		normalizedReq := c.normalizeSubscribeReq(req)
		args = append(args, normalizedReq)

		// 移除订阅信息
		delete(c.subscriptions, normalizedReq)
		c.allSubscribes.Remove(normalizedReq)
	}

	// 发送取消订阅消息
	wsReq := WsBaseReq{
		Op:   "unsubscribe",
		Args: args,
	}

	return c.sendMessage(wsReq)
}

// normalizeSubscribeReq 标准化订阅请求
func (c *BitgetMarketClient) normalizeSubscribeReq(req SubscribeReq) SubscribeReq {
	return SubscribeReq{
		InstType: strings.ToUpper(req.InstType),
		Channel:  strings.ToLower(req.Channel),
		InstId:   strings.ToUpper(req.InstId),
	}
}

// sendMessage 发送消息
func (c *BitgetMarketClient) sendMessage(msg interface{}) error {
	if c.conn == nil {
		return errors.New("WebSocket连接不可用")
	}

	c.sendMutex.Lock()
	defer c.sendMutex.Unlock()

	data, err := sonic.Marshal(msg)
	if err != nil {
		return fmt.Errorf("序列化消息失败: %v", err)
	}

	log.Info("发送消息: %s", string(data))

	c.conn.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout))
	return c.conn.WriteMessage(websocket.TextMessage, data)
}

// readLoop 消息读取循环
func (c *BitgetMarketClient) readLoop() {
	defer func() {
		if r := recover(); r != nil {
			log.Error("readLoop panic: %v", r)
		}
	}()

	for {
		select {
		case <-c.ctx.Done():
			return
		default:
			if c.conn == nil {
				time.Sleep(time.Second)
				continue
			}

			c.conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
			_, message, err := c.conn.ReadMessage()
			if err != nil {
				log.Error("读取消息失败: %v", err)
				c.handleDisconnect()
				continue
			}

			c.handleMessage(string(message))
		}
	}
}

// handleMessage 处理接收到的消息
func (c *BitgetMarketClient) handleMessage(message string) {
	log.Info("收到消息: %s", message)

	// 处理pong消息
	if message == "pong" {
		c.lastPongTime = time.Now()
		log.Debug("收到pong消息")
		return
	}

	// 解析JSON消息
	var jsonMap map[string]interface{}
	if err := json.Unmarshal([]byte(message), &jsonMap); err != nil {
		log.Error("解析JSON消息失败: %v", err)
		if c.errorListener != nil {
			c.errorListener(message)
		}
		return
	}

	// 检查错误码
	if code, exists := jsonMap["code"]; exists {
		if codeNum, ok := code.(float64); ok && int(codeNum) != 0 {
			log.Error("收到错误消息: %s", message)
			if c.errorListener != nil {
				c.errorListener(message)
			}
			return
		}
	}

	// 处理订阅确认消息
	if event, exists := jsonMap["event"]; exists && event == "subscribe" {
		log.Info("订阅确认: %s", message)
		return
	}

	// 处理数据推送
	if data, exists := jsonMap["data"]; exists && data != nil {
		if arg, argExists := jsonMap["arg"]; argExists {
			listener := c.getListener(arg)
			if listener != nil {
				listener(message)
			}
		}
		return
	}

	// 默认处理
	if c.listener != nil {
		c.listener(message)
	}
}

// getListener 获取对应的监听器
func (c *BitgetMarketClient) getListener(argJson interface{}) OnReceive {
	argMap, ok := argJson.(map[string]interface{})
	if !ok {
		return c.listener
	}

	subscribeReq := SubscribeReq{
		InstType: fmt.Sprintf("%v", argMap["instType"]),
		Channel:  fmt.Sprintf("%v", argMap["channel"]),
		InstId:   fmt.Sprintf("%v", argMap["instId"]),
	}

	if listener, exists := c.subscriptions[subscribeReq]; exists {
		return listener
	}

	return c.listener
}

// pingLoop Ping循环
func (c *BitgetMarketClient) pingLoop() {
	ticker := time.NewTicker(c.config.PingInterval)
	defer ticker.Stop()

	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			if c.connected && c.conn != nil {
				c.sendMutex.Lock()
				c.conn.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout))
				err := c.conn.WriteMessage(websocket.TextMessage, []byte("ping"))
				c.sendMutex.Unlock()

				if err != nil {
					log.Error("发送ping失败: %v", err)
					c.handleDisconnect()
				}
			}
		}
	}
}

// monitorConnection 监控连接状态
func (c *BitgetMarketClient) monitorConnection() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			if c.connected && time.Since(c.lastPongTime) > 60*time.Second {
				log.Warn("长时间未收到pong消息，触发重连")
				c.handleDisconnect()
			}
		}
	}
}

// handleDisconnect 处理断开连接
func (c *BitgetMarketClient) handleDisconnect() {
	c.reconnectMutex.Lock()
	defer c.reconnectMutex.Unlock()

	if !c.connected {
		return
	}

	log.Warn("WebSocket连接断开，准备重连")
	c.connected = false

	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}

	// 重连逻辑
	go c.reconnect()
}

// reconnect 重连
func (c *BitgetMarketClient) reconnect() {
	for c.reconnectCount < c.config.MaxReconnectTimes {
		select {
		case <-c.ctx.Done():
			return
		default:
			c.reconnectCount++
			log.Info("尝试重连 (%d/%d)", c.reconnectCount, c.config.MaxReconnectTimes)

			time.Sleep(c.config.ReconnectInterval)

			if err := c.Connect(); err != nil {
				log.Error("重连失败: %v", err)
				continue
			}

			// 重新订阅
			if err := c.resubscribe(); err != nil {
				log.Error("重新订阅失败: %v", err)
				c.handleDisconnect()
				continue
			}

			log.Info("重连成功")
			return
		}
	}

	log.Error("达到最大重连次数，停止重连")
}

// resubscribe 重新订阅
func (c *BitgetMarketClient) resubscribe() error {
	if len(c.subscriptions) == 0 {
		return nil
	}

	var args []interface{}
	for req := range c.subscriptions {
		args = append(args, req)
	}

	wsReq := WsBaseReq{
		Op:   "subscribe",
		Args: args,
	}

	return c.sendMessage(wsReq)
}

// Close 关闭客户端
func (c *BitgetMarketClient) Close() error {
	c.cancel()
	c.connected = false

	if c.conn != nil {
		return c.conn.Close()
	}

	return nil
}

// IsConnected 检查连接状态
func (c *BitgetMarketClient) IsConnected() bool {
	return c.connected
}