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

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

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

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

	// ResponseChan 响应通道
	ResponseChan chan *WebSocketMessage
}

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

	// mutex 互斥锁
	mutex sync.RWMutex

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

	// defaultTimeout 默认超时时间(秒)
	defaultTimeout int
}

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

// SetDefaultTimeout 设置默认超时时间
func (ws *WebSocketService) SetDefaultTimeout(timeout int) {
	ws.defaultTimeout = timeout
}

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

// SetResponseHandler 设置响应处理器
func (ws *WebSocketService) SetResponseHandler(deviceID string, responseChan chan *WebSocketMessage) {
	ws.mutex.Lock()
	defer ws.mutex.Unlock()

	if deviceConn, exists := ws.connections[deviceID]; exists {
		deviceConn.ResponseChan = responseChan
	}
}

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

// CommandResponse WebSocket命令响应结构体
type CommandResponse struct {
	// DeviceID 设备ID
	DeviceID string `json:"device_id,omitempty"`

	// Command 命令名称
	Command string `json:"command,omitempty"`

	// Data 响应数据
	Data interface{} `json:"data,omitempty"`

	// Status 响应状态
	Status string `json:"status,omitempty"`

	// Message 响应消息
	Message string `json:"message,omitempty"`

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

// ParseData 将响应数据解析到目标结构体
func (cr *CommandResponse) ParseData(target interface{}) error {
	dataBytes, err := json.Marshal(cr.Data)
	if err != nil {
		return err
	}
	return json.Unmarshal(dataBytes, target)
}

// CommandResult WebSocket命令执行结果
type CommandResult struct {
	// Response 响应消息
	Response *CommandResponse

	// Error 错误信息
	Error error
}

// SendCommandAndWait 发送指令并等待响应
func (ws *WebSocketService) SendCommandAndWait(deviceID, command string, data interface{}, timeout int) (*CommandResponse, error) {
	// 如果未指定超时时间，使用默认超时时间
	if timeout <= 0 {
		timeout = ws.defaultTimeout
	}

	// 创建用于接收响应的通道
	responseChan := make(chan *WebSocketMessage, 1)
	ws.SetResponseHandler(deviceID, responseChan)

	// 发送指令
	if err := ws.SendCommand(deviceID, command, data); err != nil {
		return nil, fmt.Errorf("发送指令失败: %v", err)
	}

	// 等待设备响应，设置超时
	var response *WebSocketMessage
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer cancel()

	select {
	case response = <-responseChan:
		// 成功接收到响应
		// 转换为CommandResponse结构体
		commandResponse := &CommandResponse{
			DeviceID:  response.DeviceID,
			Command:   response.Command,
			Data:      response.Data,
			Timestamp: response.Timestamp,
		}

		// 尝试提取状态和消息字段
		if responseData, ok := response.Data.(map[string]interface{}); ok {
			if status, exists := responseData["status"]; exists {
				if statusStr, ok := status.(string); ok {
					commandResponse.Status = statusStr
				}
			}

			if message, exists := responseData["message"]; exists {
				if messageStr, ok := message.(string); ok {
					commandResponse.Message = messageStr
				}
			}
		}

		return commandResponse, nil
	case <-ctx.Done():
		// 超时处理
		return nil, fmt.Errorf("等待设备响应超时")
	}
}

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

	// 处理响应消息
	if msg.Type == MessageTypeResponse {
		ws.mutex.RLock()
		if deviceConn, exists := ws.connections[deviceID]; exists && deviceConn.ResponseChan != nil {
			// 发送响应到通道
			select {
			case deviceConn.ResponseChan <- &msg:
				// 成功发送
			default:
				// 通道已满，丢弃消息
				ws.logger.Warn(fmt.Sprintf("设备 %s 的响应通道已满，丢弃响应消息", deviceID))
			}
		}
		ws.mutex.RUnlock()
	}

	// 记录消息日志
	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
}