RobotSpineAdmin/ruoyi-admin/src/main/java/com/ruoyi/websocket/controller/WebSocketController.java

package com.ruoyi.websocket.controller;

import com.ruoyi.common.core.domain.AjaxResult;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import io.swagger.annotations.ApiParam;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.Payload;
import org.springframework.messaging.handler.annotation.SendTo;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.web.bind.annotation.*;
import com.alibaba.fastjson2.JSONObject;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket控制器
 * 处理前端通过WebSocket发送的消息（如订阅设备、发送指令等）
 */
@Api(tags = "WebSocket机器人控制")
@RestController
@RequestMapping("/ws/robot")
public class WebSocketController {

    /**
     * 记录当前连接的设备ID（设备ID -> SessionID）
     */
    private final Map<String, String> deviceSessions = new ConcurrentHashMap<>();

    /**
     * 测试WebSocket连接是否正常
     */
    @ApiOperation("测试WebSocket连接")
    @GetMapping("/test")
    public AjaxResult test() {
        return AjaxResult.success("WebSocket连接正常", null);
    }

    /**
     * 客户端通过WebSocket发送消息（STOMP协议）
     * 客户端发送格式：/app/message
     * { "deviceId": "ld000001", "action": "subscribe", "data": {} }
     */
    @ApiOperation("处理WebSocket消息")
    @MessageMapping("/message")
    @SendTo("/topic/robot/response")
    public Map<String, Object> handleMessage(
            @Payload Map<String, Object> message,
            SimpMessageHeaderAccessor headerAccessor
    ) {
        Map<String, Object> response = new HashMap<>();
        String deviceId = (String) message.get("deviceId");
        String action = (String) message.get("action");
        
        response.put("success", true);
        response.put("deviceId", deviceId);
        response.put("action", action);
        response.put("timestamp", System.currentTimeMillis());
        
        // 记录设备会话
        if (deviceId != null && headerAccessor.getSessionId() != null) {
            deviceSessions.put(deviceId, headerAccessor.getSessionId());
            response.put("message", "设备 " + deviceId + " 已注册");
        }
        
        return response;
    }

    /**
     * 注册设备（HTTP接口，用于前端主动注册）
     */
    @ApiOperation("注册设备连接")
    @PostMapping("/register")
    public AjaxResult registerDevice(
            @ApiParam("设备ID") @RequestParam String deviceId,
            @ApiParam("会话ID") @RequestParam(required = false) String sessionId
    ) {
        if (sessionId != null) {
            deviceSessions.put(deviceId, sessionId);
        }
        return AjaxResult.success("设备注册成功", deviceId);
    }

    /**
     * 获取当前在线的设备列表
     */
    @ApiOperation("获取在线设备列表")
    @GetMapping("/online")
    public AjaxResult getOnlineDevices() {
        return AjaxResult.success(deviceSessions.keySet());
    }

    /**
     * 心跳检测
     */
    @ApiOperation("心跳检测")
    @PostMapping("/ping")
    public AjaxResult ping(@ApiParam("设备ID") @RequestParam String deviceId) {
        Map<String, Object> data = new HashMap<>();
        data.put("deviceId", deviceId);
        data.put("status", "online");
        data.put("timestamp", System.currentTimeMillis());
        return AjaxResult.success("心跳正常", data);
    }
}