exchange_go/services/bitgetservice/bitget_socketmanager.go

package bitgetservice

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"go-admin/pkg/utility"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/bytedance/sonic"
	log "github.com/go-admin-team/go-admin-core/logger"
	"github.com/gorilla/websocket"
	"golang.org/x/net/proxy"
)

// BitgetWebSocketManager Bitget WebSocket管理器
// 类似于BinanceWebSocketManager，支持外部传入API密钥
type BitgetWebSocketManager struct {
	ws          *websocket.Conn
	stopChannel chan struct{}
	url         string
	/* 0-现货 1-合约 */
	wsType       int
	apiKey       string
	apiSecret    string
	passphrase   string
	proxyType    string
	proxyAddress string
	reconnect    chan struct{}
	isStopped    bool       // 标记 WebSocket 是否已主动停止
	mu           sync.Mutex // 用于控制并发访问 isStopped
	cancelFunc   context.CancelFunc
	reconnecting atomic.Bool // 防止重复重连
	ConnectTime  time.Time   // 当前连接建立时间
}

// 已有连接存储
var BitgetSpotSockets = map[string]*BitgetWebSocketManager{}
var BitgetFutureSockets = map[string]*BitgetWebSocketManager{}

// NewBitgetWebSocketManager 创建新的Bitget WebSocket管理器
func NewBitgetWebSocketManager(wsType int, apiKey, apiSecret, passphrase, proxyType, proxyAddress string) *BitgetWebSocketManager {
	var wsUrl string

	switch wsType {
	case 0: // 现货
		wsUrl = "wss://ws.bitget.com/spot/v1/stream"
	case 1: // 合约
		wsUrl = "wss://ws.bitget.com/mix/v1/stream"
	default:
		wsUrl = "wss://ws.bitget.com/spot/v1/stream"
	}

	return &BitgetWebSocketManager{
		stopChannel:  make(chan struct{}, 10),
		reconnect:    make(chan struct{}, 10),
		isStopped:    false,
		url:          wsUrl,
		wsType:       wsType,
		apiKey:       apiKey,
		apiSecret:    apiSecret,
		passphrase:   passphrase,
		proxyType:    proxyType,
		proxyAddress: proxyAddress,
	}
}

// GetKey 获取API Key
func (wm *BitgetWebSocketManager) GetKey() string {
	return wm.apiKey
}

// Start 启动WebSocket连接
func (wm *BitgetWebSocketManager) Start() {
	utility.SafeGo(wm.run)
}

// Restart 重启连接
func (wm *BitgetWebSocketManager) Restart(apiKey, apiSecret, passphrase, proxyType, proxyAddress string) *BitgetWebSocketManager {
	wm.mu.Lock()
	defer wm.mu.Unlock()

	wm.apiKey = apiKey
	wm.apiSecret = apiSecret
	wm.passphrase = passphrase
	wm.proxyType = proxyType
	wm.proxyAddress = proxyAddress

	if wm.isStopped {
		wm.isStopped = false
		utility.SafeGo(wm.run)
	} else {
		log.Warnf("调用Bitget WebSocket restart")
		wm.triggerReconnect(true)
	}

	return wm
}

// triggerReconnect 触发重连
func (wm *BitgetWebSocketManager) triggerReconnect(force bool) {
	if force {
		wm.reconnecting.Store(false)
	}

	if wm.reconnecting.CompareAndSwap(false, true) {
		log.Warnf("准备重连Bitget WebSocket key: %s wsType: %v", wm.apiKey, wm.wsType)
		select {
		case wm.reconnect <- struct{}{}:
		default:
			log.Debugf("Bitget reconnect 信号已存在，跳过 key:%s", wm.apiKey)
		}
	}
}

// run 主运行循环
func (wm *BitgetWebSocketManager) run() {
	ctx, cancel := context.WithCancel(context.Background())
	wm.cancelFunc = cancel

	utility.SafeGo(wm.handleSignal)

	for {
		select {
		case <-ctx.Done():
			return
		default:
			if err := wm.connect(ctx); err != nil {
				wm.handleConnectionError(err)

				if wm.isErrorCountExceeded() {
					log.Error("连接Bitget WebSocket时出错次数过多，停止WebSocket管理器: wsType=%v, key=%s", wm.wsType, wm.apiKey)
					wm.Stop()
					return
				}

				time.Sleep(5 * time.Second)
				continue
			}

			<-wm.stopChannel
			log.Info("停止Bitget WebSocket管理器... wsType=%s", getBitgetWsTypeName(wm.wsType))
			wm.Stop()
			return
		}
	}
}

// handleConnectionError 处理连接错误
func (wm *BitgetWebSocketManager) handleConnectionError(err error) {
	// 这里可以添加错误记录到Redis的逻辑，类似binance的实现
	log.Error("连接Bitget WebSocket时出错: wsType=%v, key=%v, 错误=%v", wm.wsType, wm.apiKey, err)
}

// isErrorCountExceeded 检查错误次数是否超过阈值
func (wm *BitgetWebSocketManager) isErrorCountExceeded() bool {
	// 简化版本，可以根据需要实现更复杂的错误计数逻辑
	return false
}

// handleSignal 处理终止信号
func (wm *BitgetWebSocketManager) handleSignal() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	<-ch
	wm.Stop()
}

// connect 建立WebSocket连接
func (wm *BitgetWebSocketManager) connect(ctx context.Context) error {
	dialer, err := wm.getDialer()
	if err != nil {
		return err
	}

	// 对于私有频道，需要进行身份验证
	headers := wm.getAuthHeaders()

	wm.ws, _, err = dialer.Dial(wm.url, headers)
	if err != nil {
		return err
	}

	// 连接成功，更新连接时间
	wm.ConnectTime = time.Now()
	log.Info("已连接到Bitget %s WebSocket key:%s", getBitgetWsTypeName(wm.wsType), wm.apiKey)

	// 设置Ping处理器
	wm.ws.SetPingHandler(func(appData string) error {
		log.Info("收到Bitget Ping消息 wstype:%v key:%s data:%s", wm.wsType, wm.apiKey, appData)

		for x := 0; x < 5; x++ {
			if err := wm.ws.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second*10)); err != nil {
				log.Error("Bitget回应Pong失败 次数:%d err:%v", x, err)
				time.Sleep(time.Second * 1)
				continue
			}
			break
		}
		return nil
	})

	// 启动各种处理协程
	utility.SafeGo(func() { wm.readMessages(ctx) })
	utility.SafeGo(func() { wm.handleReconnect(ctx) })
	utility.SafeGo(func() { wm.startPingLoop(ctx) })

	// 发送登录认证消息
	if err := wm.authenticate(); err != nil {
		return fmt.Errorf("bitget WebSocket认证失败: %v", err)
	}

	return nil
}

// getDialer 获取WebSocket拨号器
func (wm *BitgetWebSocketManager) getDialer() (*websocket.Dialer, error) {
	if wm.proxyAddress == "" {
		return &websocket.Dialer{}, nil
	}

	if !strings.HasPrefix(wm.proxyAddress, "http://") &&
		!strings.HasPrefix(wm.proxyAddress, "https://") &&
		!strings.HasPrefix(wm.proxyAddress, "socks5://") {
		wm.proxyAddress = wm.proxyType + "://" + wm.proxyAddress
	}

	proxyURL, err := url.Parse(wm.proxyAddress)
	if err != nil {
		return nil, fmt.Errorf("failed to parse proxy URL: %v", err)
	}

	switch proxyURL.Scheme {
	case "socks5":
		return wm.createSocks5Dialer(proxyURL)
	case "http", "https":
		return &websocket.Dialer{
			Proxy: func(req *http.Request) (*url.URL, error) {
				return proxyURL, nil
			},
		}, nil
	default:
		return nil, fmt.Errorf("unsupported proxy scheme: %s", proxyURL.Scheme)
	}
}

// createSocks5Dialer 创建SOCKS5代理拨号器
func (wm *BitgetWebSocketManager) createSocks5Dialer(proxyURL *url.URL) (*websocket.Dialer, error) {
	auth := &proxy.Auth{}
	if proxyURL.User != nil {
		auth.User = proxyURL.User.Username()
		auth.Password, _ = proxyURL.User.Password()
	}

	socksDialer, err := proxy.SOCKS5("tcp", proxyURL.Host, auth, proxy.Direct)
	if err != nil {
		return nil, fmt.Errorf("failed to create SOCKS5 proxy dialer: %v", err)
	}

	return &websocket.Dialer{
		NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return socksDialer.Dial(network, addr)
		},
	}, nil
}

// getAuthHeaders 获取认证头信息
func (wm *BitgetWebSocketManager) getAuthHeaders() http.Header {
	headers := http.Header{}

	if wm.apiKey != "" {
		timestamp := strconv.FormatInt(time.Now().Unix(), 10)
		sign := wm.generateSignature(timestamp)

		headers.Set("ACCESS-KEY", wm.apiKey)
		headers.Set("ACCESS-SIGN", sign)
		headers.Set("ACCESS-TIMESTAMP", timestamp)
		headers.Set("ACCESS-PASSPHRASE", wm.passphrase)
	}

	return headers
}

// generateSignature 生成签名
func (wm *BitgetWebSocketManager) generateSignature(timestamp string) string {
	message := timestamp + "GET" + "/user/verify"

	h := hmac.New(sha256.New, []byte(wm.apiSecret))
	h.Write([]byte(message))

	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

// authenticate 发送认证消息
func (wm *BitgetWebSocketManager) authenticate() error {
	if wm.apiKey == "" {
		return nil // 公共频道不需要认证
	}

	timestamp := strconv.FormatInt(time.Now().Unix(), 10)
	sign := wm.generateSignature(timestamp)

	authMsg := map[string]interface{}{
		"op": "login",
		"args": []map[string]string{
			{
				"apiKey":     wm.apiKey,
				"passphrase": wm.passphrase,
				"timestamp":  timestamp,
				"sign":       sign,
			},
		},
	}

	msgBytes, err := sonic.Marshal(authMsg)
	if err != nil {
		return err
	}

	return wm.ws.WriteMessage(websocket.TextMessage, msgBytes)
}

// readMessages 读取消息
func (wm *BitgetWebSocketManager) readMessages(ctx context.Context) {
	defer func() {
		if r := recover(); r != nil {
			log.Error("Bitget WebSocket消息处理器panic: %v", r)
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			wm.mu.Lock()
			if wm.isStopped || wm.ws == nil {
				wm.mu.Unlock()
				return
			}

			_, message, err := wm.ws.ReadMessage()
			wm.mu.Unlock()

			if err != nil {
				log.Error("Bitget WebSocket读取消息错误: %v", err)
				wm.triggerReconnect(false)
				return
			}

			// 处理消息
			wm.handleMessage(message)
		}
	}
}

// handleMessage 处理收到的消息
func (wm *BitgetWebSocketManager) handleMessage(message []byte) {
	// 解析消息并处理
	var baseMsg map[string]interface{}
	if err := sonic.Unmarshal(message, &baseMsg); err != nil {
		log.Error("解析Bitget WebSocket消息失败: %v", err)
		return
	}

	// 检查是否是认证响应
	if op, exists := baseMsg["op"]; exists && op == "login" {
		wm.handleAuthResponse(baseMsg)
		return
	}

	// 检查是否是订阅数据
	if arg, exists := baseMsg["arg"]; exists {
		wm.handleSubscriptionData(baseMsg, arg)
		return
	}

	log.Info("收到Bitget WebSocket消息: %s", string(message))
}

// handleAuthResponse 处理认证响应
func (wm *BitgetWebSocketManager) handleAuthResponse(msg map[string]interface{}) {
	if code, exists := msg["code"]; exists {
		if code == "0" || code == 0 {
			log.Info("Bitget WebSocket认证成功 key:%s", wm.apiKey)
			// 认证成功后可以订阅私有频道
			wm.subscribePrivateChannels()
		} else {
			log.Error("Bitget WebSocket认证失败 key:%s code:%v", wm.apiKey, code)
		}
	}
}

// handleSubscriptionData 处理订阅数据
func (wm *BitgetWebSocketManager) handleSubscriptionData(msg map[string]interface{}, arg interface{}) {
	// 根据频道类型处理不同的数据
	// 这里可以根据具体需求实现不同频道的数据处理逻辑
	log.Info("收到Bitget订阅数据: %v", arg)
}

// subscribePrivateChannels 订阅私有频道
func (wm *BitgetWebSocketManager) subscribePrivateChannels() {
	// 订阅账户更新
	wm.subscribeAccount()
	// 订阅订单更新
	wm.subscribeOrders()
	// 订阅持仓更新
	wm.subscribePositions()
}

// subscribeAccount 订阅账户更新
func (wm *BitgetWebSocketManager) subscribeAccount() {
	var channel string
	if wm.wsType == 0 {
		channel = "account"
	} else {
		channel = "account"
	}

	subMsg := map[string]interface{}{
		"op": "subscribe",
		"args": []map[string]string{
			{
				"instType": wm.getInstType(),
				"channel":  channel,
				"instId":   "default",
			},
		},
	}

	wm.sendMessage(subMsg)
}

// subscribeOrders 订阅订单更新
func (wm *BitgetWebSocketManager) subscribeOrders() {
	var channel string
	if wm.wsType == 0 {
		channel = "orders"
	} else {
		channel = "orders"
	}

	subMsg := map[string]interface{}{
		"op": "subscribe",
		"args": []map[string]string{
			{
				"instType": wm.getInstType(),
				"channel":  channel,
				"instId":   "default",
			},
		},
	}

	wm.sendMessage(subMsg)
}

// subscribePositions 订阅持仓更新（仅合约）
func (wm *BitgetWebSocketManager) subscribePositions() {
	if wm.wsType != 1 {
		return // 只有合约才有持仓
	}

	subMsg := map[string]interface{}{
		"op": "subscribe",
		"args": []map[string]string{
			{
				"instType": "umcbl",
				"channel":  "positions",
				"instId":   "default",
			},
		},
	}

	wm.sendMessage(subMsg)
}

// getInstType 获取产品类型
func (wm *BitgetWebSocketManager) getInstType() string {
	if wm.wsType == 0 {
		return "sp" // 现货
	}
	return "umcbl" // U本位合约
}

// SendMessage 发送消息（公开方法）
func (wm *BitgetWebSocketManager) SendMessage(msg interface{}) error {
	msgBytes, err := sonic.Marshal(msg)
	if err != nil {
		return err
	}

	wm.mu.Lock()
	defer wm.mu.Unlock()

	if wm.ws == nil {
		return errors.New("WebSocket连接不可用")
	}

	return wm.ws.WriteMessage(websocket.TextMessage, msgBytes)
}

// sendMessage 发送消息（私有方法，保持向后兼容）
func (wm *BitgetWebSocketManager) sendMessage(msg interface{}) error {
	return wm.SendMessage(msg)
}

// handleReconnect 处理重连
func (wm *BitgetWebSocketManager) handleReconnect(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-wm.reconnect:
			wm.reconnecting.Store(false)
			log.Info("开始重连Bitget WebSocket key:%s", wm.apiKey)

			wm.mu.Lock()
			if wm.ws != nil {
				wm.ws.Close()
				wm.ws = nil
			}
			wm.mu.Unlock()

			// 触发主循环重新连接
			select {
			case wm.stopChannel <- struct{}{}:
			default:
			}
		}
	}
}

// startPingLoop 启动心跳循环
func (wm *BitgetWebSocketManager) startPingLoop(ctx context.Context) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			wm.sendPing()
		}
	}
}

// sendPing 发送心跳
func (wm *BitgetWebSocketManager) sendPing() {
	pingMsg := "ping"

	wm.mu.Lock()
	defer wm.mu.Unlock()

	if wm.ws == nil {
		return
	}

	if err := wm.ws.WriteMessage(websocket.TextMessage, []byte(pingMsg)); err != nil {
		log.Error("Bitget WebSocket发送ping失败: %v", err)
		wm.triggerReconnect(false)
	}
}

// Stop 停止WebSocket连接
func (wm *BitgetWebSocketManager) Stop() {
	wm.mu.Lock()
	defer wm.mu.Unlock()

	if wm.isStopped {
		return
	}

	wm.isStopped = true

	if wm.cancelFunc != nil {
		wm.cancelFunc()
	}

	if wm.ws != nil {
		wm.ws.Close()
		wm.ws = nil
	}

	// 从全局map中移除
	if wm.wsType == 0 {
		delete(BitgetSpotSockets, wm.apiKey)
	} else {
		delete(BitgetFutureSockets, wm.apiKey)
	}

	log.Info("Bitget WebSocket已停止 key:%s", wm.apiKey)
}

// getBitgetWsTypeName 获取WebSocket类型名称
func getBitgetWsTypeName(wsType int) string {
	switch wsType {
	case 0:
		return "现货"
	case 1:
		return "合约"
	default:
		return "未知"
	}
}

// 工具函数：获取或创建Bitget WebSocket管理器
func GetOrCreateBitgetWebSocketManager(wsType int, apiKey, apiSecret, passphrase, proxyType, proxyAddress string) *BitgetWebSocketManager {
	var socketsMap map[string]*BitgetWebSocketManager

	if wsType == 0 {
		socketsMap = BitgetSpotSockets
	} else {
		socketsMap = BitgetFutureSockets
	}

	if manager, exists := socketsMap[apiKey]; exists {
		// 如果参数有变化，重启连接
		if manager.apiSecret != apiSecret || manager.passphrase != passphrase ||
			manager.proxyType != proxyType || manager.proxyAddress != proxyAddress {
			manager.Restart(apiKey, apiSecret, passphrase, proxyType, proxyAddress)
		}
		return manager
	}

	// 创建新的管理器
	manager := NewBitgetWebSocketManager(wsType, apiKey, apiSecret, passphrase, proxyType, proxyAddress)
	socketsMap[apiKey] = manager

	return manager
}

// StopBitgetWebSocketManager 停止指定的WebSocket管理器
func StopBitgetWebSocketManager(wsType int, apiKey string) {
	var socketsMap map[string]*BitgetWebSocketManager

	if wsType == 0 {
		socketsMap = BitgetSpotSockets
	} else {
		socketsMap = BitgetFutureSockets
	}

	if manager, exists := socketsMap[apiKey]; exists {
		manager.Stop()
	}
}