┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐
│   業(yè)務(wù)系統(tǒng)API   │───>│   消息路由模塊   │───>│   連接管理模塊   │───> 客戶端
└─────────────────┘    └─────────────────┘    └─────────────────┘
        │                       │                       │
        ▼                       ▼                       ▼
┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐
│  離線消息存儲    │<───│  集群通信模塊    │<───│  會話管理模塊    │
└─────────────────┘    └─────────────────┘    └─────────────────┘

<!-- WebSocket核心依賴 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>
<!-- Redis（集群通信與會話存儲） -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private MessageHandler messageHandler; // 自定義消息處理器

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 配置WebSocket端點(diǎn)：客戶端通過ws://ip:port/ws?token=xxx連接
        registry.addHandler(messageHandler, "/ws")
                .setAllowedOrigins("*") // 允許跨域（生產(chǎn)環(huán)境需限制域名）
                .addInterceptors(new HandshakeInterceptor() {
                    // 握手前驗(yàn)證用戶（如Token解析用戶ID）
                    @Override
                    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                                  WebSocketHandler wsHandler, Map<String, Object> attributes) {
                        String token = ((ServletServerHttpRequest) request).getServletRequest().getParameter("token");
                        String userId = parseToken(token); // 解析Token獲取用戶ID
                        if (userId == null) {
                            response.setStatusCode(HttpStatus.UNAUTHORIZED);
                            return false; // 驗(yàn)證失敗，拒絕連接
                        }
                        attributes.put("userId", userId); // 存儲用戶ID到屬性中
                        return true;
                    }

                    @Override
                    public void afterHandshake(...) {}
                });
    }
}

@Component
public class MessageHandler extends TextWebSocketHandler {
    // 本地會話映射：用戶ID -> WebSocket會話（單機(jī)用）
    private final Map<String, WebSocketSession> localSessions = new ConcurrentHashMap<>();
    @Autowired
    private StringRedisTemplate redisTemplate; // Redis操作模板
    @Autowired
    private OfflineMessageService offlineMessageService; // 離線消息服務(wù)

    // 連接建立時：綁定用戶ID與Session
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String userId = (String) session.getAttributes().get("userId");
        localSessions.put(userId, session);
        
        // 1. 拉取離線消息并推送
        List<Message> offlineMessages = offlineMessageService.getByUserId(userId);
        for (Message msg : offlineMessages) {
            sendMessage(session, msg.getContent());
        }
        offlineMessageService.deleteByUserId(userId); // 清除已推送的離線消息
        
        // 2. 集群：將用戶ID與當(dāng)前節(jié)點(diǎn)IP存入Redis（便于其他節(jié)點(diǎn)定位）
        redisTemplate.opsForValue().set("user:session:" + userId, getLocalIp());
    }

    // 連接關(guān)閉時：移除會話映射
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String userId = (String) session.getAttributes().get("userId");
        localSessions.remove(userId);
        redisTemplate.delete("user:session:" + userId); // 集群：刪除Redis映射
    }

    // 發(fā)送消息到指定用戶（核心方法）
    public void pushToUser(String userId, String content) {
        // 1. 先檢查本地是否有該用戶的連接
        WebSocketSession session = localSessions.get(userId);
        if (session != null && session.isOpen()) {
            sendMessage(session, content);
            return;
        }

        // 2. 集群：檢查用戶是否連接在其他節(jié)點(diǎn)
        String targetNodeIp = redisTemplate.opsForValue().get("user:session:" + userId);
        if (targetNodeIp != null) {
            // 發(fā)布消息到Redis頻道，目標(biāo)節(jié)點(diǎn)訂閱后發(fā)送
            redisTemplate.convertAndSend("push:channel", 
                JSON.toJSONString(new PushMessage(userId, content)));
            return;
        }

        // 3. 用戶不在線：存入離線消息
        offlineMessageService.save(new Message(userId, content, LocalDateTime.now()));
    }

    // 實(shí)際發(fā)送消息（封裝異常處理）
    private void sendMessage(WebSocketSession session, String content) {
        try {
            session.sendMessage(new TextMessage(content));
        } catch (IOException e) {
            log.error("消息發(fā)送失敗", e);
        }
    }
}

@Component
public class RedisMessageListener {
    @Autowired
    private MessageHandler messageHandler;

    @Bean
    public RedisMessageListenerContainer container(RedisConnectionFactory connectionFactory) {
        RedisMessageListenerContainer container = new RedisMessageListenerContainer();
        container.setConnectionFactory(connectionFactory);
        // 訂閱推送頻道，接收其他節(jié)點(diǎn)的消息
        container.addMessageListener((message, pattern) -> {
            String json = new String(message.getBody());
            PushMessage pushMsg = JSON.parseObject(json, PushMessage.class);
            // 調(diào)用本地消息處理器發(fā)送（此時用戶應(yīng)在當(dāng)前節(jié)點(diǎn)）
            messageHandler.pushToUser(pushMsg.getUserId(), pushMsg.getContent());
        }, new ChannelTopic("push:channel"));
        return container;
    }
}

@RestController
@RequestMapping("/api/message")
public class MessageController {
    @Autowired
    private MessageHandler messageHandler;

    // 業(yè)務(wù)系統(tǒng)調(diào)用此接口觸發(fā)推送
    @PostMapping("/push")
    public Result push(@RequestBody PushRequest request) {
        messageHandler.pushToUser(request.getUserId(), request.getContent());
        return Result.success();
    }
}