pig-farm-controller/internal/api/websocket.go

// Package api 提供统一的API接口层
// 负责处理所有外部请求，包括HTTP和WebSocket接口
// 将请求路由到相应的服务层进行处理
package api

import (
	"net/http"
	"sync"
	"time"

	"git.huangwc.com/pig/pig-farm-controller/internal/logs"
	"git.huangwc.com/pig/pig-farm-controller/internal/service"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// WebSocket消息类型常量
const (
	// MessageTypeCommand 平台向设备发送的指令
	MessageTypeCommand = "command"

	// MessageTypeResponse 设备向平台发送的响应
	MessageTypeResponse = "response"

	// MessageTypeHeartbeat 心跳消息
	MessageTypeHeartbeat = "heartbeat"
)

// WebSocketMessage WebSocket消息结构
type WebSocketMessage struct {
	// Type 消息类型
	Type string `json:"type"`

	// DeviceID 设备ID
	DeviceID string `json:"device_id,omitempty"`

	// Command 指令内容
	Command string `json:"command,omitempty"`

	// Data 消息数据
	Data interface{} `json:"data,omitempty"`

	// Timestamp 时间戳
	Timestamp time.Time `json:"timestamp"`
}

// WebSocketManager WebSocket管理器
type WebSocketManager struct {
	// websocketService WebSocket服务
	websocketService *service.WebSocketService

	// logger 日志记录器
	logger *logs.Logger

	// upgrader WebSocket升级器
	upgrader websocket.Upgrader

	// mutex 互斥锁
	mutex sync.RWMutex

	// connections 设备连接映射
	connections map[string]*websocket.Conn
}

// NewWebSocketManager 创建WebSocket管理器实例
func NewWebSocketManager(websocketService *service.WebSocketService) *WebSocketManager {
	return &WebSocketManager{
		websocketService: websocketService,
		logger:           logs.NewLogger(),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				// 允许所有跨域请求
				return true
			},
		},
		connections: make(map[string]*websocket.Conn),
	}
}

// HandleConnection 处理WebSocket连接
func (wm *WebSocketManager) HandleConnection(c *gin.Context) {
	// 升级HTTP连接到WebSocket
	conn, err := wm.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		wm.logger.Error("WebSocket连接升级失败: " + err.Error())
		return
	}

	// 获取设备ID
	deviceID := c.Query("device_id")
	if deviceID == "" {
		wm.logger.Error("缺少设备ID参数")
		conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "缺少设备ID参数"))
		conn.Close()
		return
	}

	// 添加连接到映射
	wm.mutex.Lock()
	wm.connections[deviceID] = conn
	wm.mutex.Unlock()

	wm.logger.Info("设备 " + deviceID + " 已连接")

	// 发送连接成功消息
	successMsg := service.WebSocketMessage{
		Type:      "system",
		Command:   "connected",
		Timestamp: time.Now(),
	}
	conn.WriteJSON(successMsg)

	// 处理消息循环
	for {
		// 读取消息
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			wm.logger.Error("读取设备 " + deviceID + " 消息失败: " + err.Error())
			break
		}

		// 只处理文本消息
		if messageType != websocket.TextMessage {
			continue
		}

		// 处理设备消息
		if err := wm.websocketService.HandleMessage(deviceID, message); err != nil {
			wm.logger.Error("处理设备 " + deviceID + " 消息失败: " + err.Error())
			continue
		}
	}

	// 连接断开时清理
	wm.mutex.Lock()
	delete(wm.connections, deviceID)
	wm.mutex.Unlock()

	conn.Close()
	wm.logger.Info("设备 " + deviceID + " 已断开连接")
}

// SendCommand 向指定设备发送指令
func (wm *WebSocketManager) SendCommand(deviceID, command string, data interface{}) error {
	wm.mutex.RLock()
	conn, exists := wm.connections[deviceID]
	wm.mutex.RUnlock()

	if !exists {
		return wm.websocketService.SendCommand(deviceID, command, data)
	}

	// 构造消息
	msg := service.WebSocketMessage{
		Type:      service.MessageTypeCommand,
		Command:   command,
		Data:      data,
		Timestamp: time.Now(),
	}

	// 发送消息
	if err := conn.WriteJSON(msg); err != nil {
		return err
	}

	return nil
}

// GetConnectedDevices 获取已连接的设备列表
func (wm *WebSocketManager) GetConnectedDevices() []string {
	wm.mutex.RLock()
	defer wm.mutex.RUnlock()

	devices := make([]string, 0, len(wm.connections))
	for deviceID := range wm.connections {
		devices = append(devices, deviceID)
	}

	return devices
}