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

// Package core 提供WebSocket服务功能
// 实现中继设备和平台之间的双向通信
package core

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"git.huangwc.com/pig/pig-farm-controller/internal/logs"
	"git.huangwc.com/pig/pig-farm-controller/internal/model"
	"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"`
}

// DeviceConnection 设备连接信息
type DeviceConnection struct {
	// DeviceID 设备ID
	DeviceID string

	// Connection WebSocket连接
	Connection *websocket.Conn

	// LastHeartbeat 最后心跳时间
	LastHeartbeat time.Time

	// DeviceInfo 设备信息
	DeviceInfo *model.Device
}

// WebSocketService WebSocket服务
type WebSocketService struct {
	// connections 设备连接映射
	connections map[string]*DeviceConnection

	// mutex 互斥锁
	mutex sync.RWMutex

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

// NewWebSocketService 创建WebSocket服务实例
func NewWebSocketService() *WebSocketService {
	return &WebSocketService{
		connections: make(map[string]*DeviceConnection),
		logger:      logs.NewLogger(),
	}
}

// AddConnection 添加设备连接
func (ws *WebSocketService) AddConnection(deviceID string, conn *websocket.Conn) {
	ws.mutex.Lock()
	defer ws.mutex.Unlock()

	ws.connections[deviceID] = &DeviceConnection{
		DeviceID:      deviceID,
		Connection:    conn,
		LastHeartbeat: time.Now(),
	}

	ws.logger.Info(fmt.Sprintf("设备 %s 已连接", deviceID))
}

// RemoveConnection 移除设备连接
func (ws *WebSocketService) RemoveConnection(deviceID string) {
	ws.mutex.Lock()
	defer ws.mutex.Unlock()

	delete(ws.connections, deviceID)

	ws.logger.Info(fmt.Sprintf("设备 %s 已断开连接", deviceID))
}

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

	if !exists {
		return fmt.Errorf("设备 %s 未连接", deviceID)
	}

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

	// 发送消息
	if err := deviceConn.Connection.WriteJSON(msg); err != nil {
		return fmt.Errorf("向设备 %s 发送指令失败: %v", deviceID, err)
	}

	return nil
}

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

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

	return devices
}

// HandleMessage 处理来自设备的消息
func (ws *WebSocketService) HandleMessage(deviceID string, message []byte) error {
	// 解析消息
	var msg WebSocketMessage
	if err := json.Unmarshal(message, &msg); err != nil {
		return fmt.Errorf("解析设备 %s 消息失败: %v", deviceID, err)
	}

	// 更新心跳时间
	if msg.Type == MessageTypeHeartbeat {
		ws.mutex.Lock()
		if deviceConn, exists := ws.connections[deviceID]; exists {
			deviceConn.LastHeartbeat = time.Now()
		}
		ws.mutex.Unlock()
	}

	// 记录消息日志
	ws.logger.Info(fmt.Sprintf("收到来自设备 %s 的消息: %v", deviceID, msg))

	return nil
}

// GetDeviceConnection 获取设备连接信息
func (ws *WebSocketService) GetDeviceConnection(deviceID string) (*DeviceConnection, bool) {
	ws.mutex.RLock()
	defer ws.mutex.RUnlock()

	deviceConn, exists := ws.connections[deviceID]
	return deviceConn, exists
}