package com.example.websocket.config;

import com.example.websocket.handler.*;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.*;

/**
 * Spring Boot 4 - WebSocket 配置
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    
    private final ChatWebSocketHandler chatHandler;
    private final NotificationWebSocketHandler notificationHandler;
    
    public WebSocketConfig(
            ChatWebSocketHandler chatHandler,
            NotificationWebSocketHandler notificationHandler) {
        this.chatHandler = chatHandler;
        this.notificationHandler = notificationHandler;
    }
    
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler, "/ws/chat")
            .setAllowedOrigins("*");
        
        registry.addHandler(notificationHandler, "/ws/notifications")
            .setAllowedOrigins("*");
    }
}

spring:
  application:
    name: websocket-demo
  threads:
    virtual:
      enabled: true  # WebSocket 自动使用虚拟线程

server:
  port: 8080

logging:
  level:
    com.example.websocket: DEBUG
    org.springframework.web.socket: DEBUG

package com.example.websocket.model;

import java.time.Instant;

/**
 * 聊天消息 - 使用 Record
 */
public record ChatMessage(
    String id,
    String sender,
    String content,
    String room,
    MessageType type,
    Instant timestamp
) {
    public enum MessageType {
        CHAT, JOIN, LEAVE, TYPING
    }
    
    public static ChatMessage chat(String sender, String content, String room) {
        return new ChatMessage(
            java.util.UUID.randomUUID().toString(),
            sender,
            content,
            room,
            MessageType.CHAT,
            Instant.now()
        );
    }
    
    public static ChatMessage join(String sender, String room) {
        return new ChatMessage(
            java.util.UUID.randomUUID().toString(),
            sender,
            sender + " joined the room",
            room,
            MessageType.JOIN,
            Instant.now()
        );
    }
    
    public static ChatMessage leave(String sender, String room) {
        return new ChatMessage(
            java.util.UUID.randomUUID().toString(),
            sender,
            sender + " left the room",
            room,
            MessageType.LEAVE,
            Instant.now()
        );
    }
}

package com.example.websocket.model;

import java.time.Instant;

public record Notification(
    String id,
    String userId,
    String title,
    String message,
    NotificationType type,
    Instant timestamp
) {
    public enum NotificationType {
        INFO, WARNING, ERROR, SUCCESS
    }
    
    public static Notification info(String userId, String title, String message) {
        return new Notification(
            java.util.UUID.randomUUID().toString(),
            userId,
            title,
            message,
            NotificationType.INFO,
            Instant.now()
        );
    }
}

package com.example.websocket.handler;

import com.example.websocket.model.ChatMessage;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * Spring Boot 4 - 聊天 WebSocket 处理器
 * 使用虚拟线程处理连接
 */
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(ChatWebSocketHandler.class);
    
    private final ObjectMapper objectMapper;
    
    // 房间 -> 会话集合
    private final Map<String, CopyOnWriteArraySet<WebSocketSession>> rooms = 
        new ConcurrentHashMap<>();
    
    // 会话 -> 用户信息
    private final Map<String, UserInfo> sessionUsers = new ConcurrentHashMap<>();
    
    public ChatWebSocketHandler(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = getUserId(session);
        String room = getRoom(session);
        
        log.info("WebSocket connected: userId={}, room={}, sessionId={}, thread={}", 
            userId, room, session.getId(), Thread.currentThread());
        
        // 添加到房间
        rooms.computeIfAbsent(room, k -> new CopyOnWriteArraySet<>()).add(session);
        sessionUsers.put(session.getId(), new UserInfo(userId, room));
        
        // 发送加入消息
        ChatMessage joinMessage = ChatMessage.join(userId, room);
        broadcastToRoom(room, joinMessage, session.getId());
        
        // 发送当前在线用户列表
        sendOnlineUsers(session, room);
    }
    
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) 
            throws Exception {
        
        log.debug("Received message: {}, thread={}", 
            message.getPayload(), Thread.currentThread());
        
        ChatMessage chatMessage = objectMapper.readValue(
            message.getPayload(), 
            ChatMessage.class
        );
        
        String room = sessionUsers.get(session.getId()).room();
        
        // 广播消息到房间
        broadcastToRoom(room, chatMessage, null);
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) 
            throws Exception {
        
        UserInfo userInfo = sessionUsers.remove(session.getId());
        if (userInfo != null) {
            String room = userInfo.room();
            
            // 从房间移除
            CopyOnWriteArraySet<WebSocketSession> roomSessions = rooms.get(room);
            if (roomSessions != null) {
                roomSessions.remove(session);
                if (roomSessions.isEmpty()) {
                    rooms.remove(room);
                }
            }
            
            // 发送离开消息
            ChatMessage leaveMessage = ChatMessage.leave(userInfo.userId(), room);
            broadcastToRoom(room, leaveMessage, null);
            
            log.info("WebSocket disconnected: userId={}, room={}, status={}", 
                userInfo.userId(), room, status);
        }
    }
    
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) 
            throws Exception {
        log.error("WebSocket error: sessionId={}", session.getId(), exception);
    }
    
    /**
     * 广播消息到房间
     */
    private void broadcastToRoom(String room, ChatMessage message, String excludeSessionId) {
        CopyOnWriteArraySet<WebSocketSession> sessions = rooms.get(room);
        if (sessions == null) {
            return;
        }
        
        String messageJson;
        try {
            messageJson = objectMapper.writeValueAsString(message);
        } catch (Exception e) {
            log.error("Failed to serialize message", e);
            return;
        }
        
        TextMessage textMessage = new TextMessage(messageJson);
        
        // 使用虚拟线程并行发送
        sessions.parallelStream()
            .filter(s -> !s.getId().equals(excludeSessionId))
            .filter(WebSocketSession::isOpen)
            .forEach(s -> {
                try {
                    synchronized (s) {
                        s.sendMessage(textMessage);
                    }
                } catch (IOException e) {
                    log.error("Failed to send message to session: {}", s.getId(), e);
                }
            });
    }
    
    /**
     * 发送在线用户列表
     */
    private void sendOnlineUsers(WebSocketSession session, String room) {
        CopyOnWriteArraySet<WebSocketSession> sessions = rooms.get(room);
        if (sessions == null) {
            return;
        }
        
        var onlineUsers = sessions.stream()
            .map(s -> sessionUsers.get(s.getId()))
            .filter(u -> u != null)
            .map(UserInfo::userId)
            .toList();
        
        try {
            String json = objectMapper.writeValueAsString(Map.of(
                "type", "ONLINE_USERS",
                "users", onlineUsers
            ));
            session.sendMessage(new TextMessage(json));
        } catch (IOException e) {
            log.error("Failed to send online users", e);
        }
    }
    
    private String getUserId(WebSocketSession session) {
        // 从查询参数获取用户ID
        String query = session.getUri().getQuery();
        if (query != null && query.contains("userId=")) {
            return query.split("userId=")[1].split("&")[0];
        }
        return "anonymous";
    }
    
    private String getRoom(WebSocketSession session) {
        String query = session.getUri().getQuery();
        if (query != null && query.contains("room=")) {
            return query.split("room=")[1].split("&")[0];
        }
        return "default";
    }
    
    record UserInfo(String userId, String room) {}
}

package com.example.websocket.handler;

import com.example.websocket.model.Notification;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 通知 WebSocket 处理器
 */
@Component
public class NotificationWebSocketHandler extends TextWebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(NotificationWebSocketHandler.class);
    
    private final ObjectMapper objectMapper;
    
    // userId -> WebSocketSession
    private final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();
    
    public NotificationWebSocketHandler(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = getUserId(session);
        userSessions.put(userId, session);
        
        log.info("Notification WebSocket connected: userId={}, thread={}", 
            userId, Thread.currentThread());
        
        // 发送欢迎消息
        Notification welcome = Notification.info(
            userId, 
            "Welcome", 
            "You are now connected to notifications"
        );
        sendNotification(session, welcome);
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) 
            throws Exception {
        String userId = getUserId(session);
        userSessions.remove(userId);
        
        log.info("Notification WebSocket disconnected: userId={}", userId);
    }
    
    /**
     * 发送通知给指定用户
     */
    public void sendNotificationToUser(String userId, Notification notification) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            sendNotification(session, notification);
        }
    }
    
    /**
     * 广播通知给所有用户
     */
    public void broadcastNotification(Notification notification) {
        userSessions.values().parallelStream()
            .filter(WebSocketSession::isOpen)
            .forEach(session -> sendNotification(session, notification));
    }
    
    private void sendNotification(WebSocketSession session, Notification notification) {
        try {
            String json = objectMapper.writeValueAsString(notification);
            synchronized (session) {
                session.sendMessage(new TextMessage(json));
            }
        } catch (IOException e) {
            log.error("Failed to send notification", e);
        }
    }
    
    private String getUserId(WebSocketSession session) {
        String query = session.getUri().getQuery();
        if (query != null && query.contains("userId=")) {
            return query.split("userId=")[1].split("&")[0];
        }
        return "anonymous";
    }
}

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Chat - Spring Boot 4</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 50px auto;
            padding: 20px;
        }
        #messages {
            border: 1px solid #ccc;
            height: 400px;
            overflow-y: scroll;
            padding: 10px;
            margin-bottom: 10px;
        }
        .message {
            margin: 5px 0;
            padding: 5px;
        }
        .join { color: green; }
        .leave { color: red; }
        .chat { color: black; }
        input[type="text"] {
            width: 70%;
            padding: 10px;
        }
        button {
            padding: 10px 20px;
        }
    </style>
</head>
<body>
    <h1>WebSocket Chat</h1>
    
    <div>
        <input type="text" id="userId" placeholder="Your name" />
        <input type="text" id="room" placeholder="Room name" value="general" />
        <button onclick="connect()">Connect</button>
        <button onclick="disconnect()">Disconnect</button>
    </div>
    
    <div id="messages"></div>
    
    <div>
        <input type="text" id="messageInput" placeholder="Type a message..." />
        <button onclick="sendMessage()">Send</button>
    </div>
    
    <script>
        let ws = null;
        
        function connect() {
            const userId = document.getElementById('userId').value;
            const room = document.getElementById('room').value;
            
            if (!userId) {
                alert('Please enter your name');
                return;
            }
            
            const url = `ws://localhost:8080/ws/chat?userId=${userId}&room=${room}`;
            ws = new WebSocket(url);
            
            ws.onopen = () => {
                addMessage('Connected to chat room: ' + room, 'join');
            };
            
            ws.onmessage = (event) => {
                const data = JSON.parse(event.data);
                
                if (data.type === 'ONLINE_USERS') {
                    console.log('Online users:', data.users);
                } else {
                    const className = data.type.toLowerCase();
                    addMessage(`[${data.sender}] ${data.content}`, className);
                }
            };
            
            ws.onclose = () => {
                addMessage('Disconnected from chat', 'leave');
            };
            
            ws.onerror = (error) => {
                console.error('WebSocket error:', error);
            };
        }
        
        function disconnect() {
            if (ws) {
                ws.close();
                ws = null;
            }
        }
        
        function sendMessage() {
            if (!ws || ws.readyState !== WebSocket.OPEN) {
                alert('Not connected');
                return;
            }
            
            const input = document.getElementById('messageInput');
            const message = input.value.trim();
            
            if (!message) return;
            
            const userId = document.getElementById('userId').value;
            const room = document.getElementById('room').value;
            
            const chatMessage = {
                sender: userId,
                content: message,
                room: room,
                type: 'CHAT'
            };
            
            ws.send(JSON.stringify(chatMessage));
            input.value = '';
        }
        
        function addMessage(text, className) {
            const messagesDiv = document.getElementById('messages');
            const messageDiv = document.createElement('div');
            messageDiv.className = 'message ' + className;
            messageDiv.textContent = text;
            messagesDiv.appendChild(messageDiv);
            messagesDiv.scrollTop = messagesDiv.scrollHeight;
        }
        
        // Enter key to send
        document.getElementById('messageInput').addEventListener('keypress', (e) => {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });
    </script>
</body>
</html>

package com.example.websocket.controller;

import com.example.websocket.model.Notification;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.time.Instant;
import java.util.Map;
import java.util.concurrent.*;

/**
 * Spring Boot 4 - SSE 控制器
 * 使用虚拟线程处理长连接
 */
@RestController
@RequestMapping("/api/sse")
public class SSEController {
    private static final Logger log = LoggerFactory.getLogger(SSEController.class);
    
    // 存储所有活跃的 SSE 连接
    private final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();
    
    /**
     * 订阅实时通知
     */
    @GetMapping(value = "/notifications", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamNotifications(@RequestParam String userId) {
        log.info("New SSE connection: userId={}, thread={}", 
            userId, Thread.currentThread());
        
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
        emitters.put(userId, emitter);
        
        // 连接建立时发送欢迎消息
        try {
            emitter.send(SseEmitter.event()
                .name("connected")
                .data(Map.of(
                    "message", "Connected to notification stream",
                    "timestamp", Instant.now()
                ))
            );
        } catch (IOException e) {
            log.error("Failed to send welcome message", e);
        }
        
        // 设置超时和完成回调
        emitter.onTimeout(() -> {
            log.info("SSE timeout: userId={}", userId);
            emitters.remove(userId);
        });
        
        emitter.onCompletion(() -> {
            log.info("SSE completed: userId={}", userId);
            emitters.remove(userId);
        });
        
        emitter.onError((ex) -> {
            log.error("SSE error: userId={}", userId, ex);
            emitters.remove(userId);
        });
        
        return emitter;
    }
    
    /**
     * 订阅实时数据（模拟股票价格）
     */
    @GetMapping(value = "/stock-prices", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamStockPrices(@RequestParam String symbol) {
        SseEmitter emitter = new SseEmitter();
        
        // 使用虚拟线程定期发送数据
        Thread.startVirtualThread(() -> {
            try {
                for (int i = 0; i < 100; i++) {
                    if (emitter.getTimeout() == 0) {
                        break;
                    }
                    
                    double price = 100 + Math.random() * 10;
                    
                    emitter.send(SseEmitter.event()
                        .name("price-update")
                        .data(Map.of(
                            "symbol", symbol,
                            "price", String.format("%.2f", price),
                            "timestamp", Instant.now()
                        ))
                    );
                    
                    Thread.sleep(1000); // 每秒更新一次
                }
                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });
        
        return emitter;
    }
    
    /**
     * 订阅服务器状态
     */
    @GetMapping(value = "/server-status", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamServerStatus() {
        SseEmitter emitter = new SseEmitter();
        
        Thread.startVirtualThread(() -> {
            try {
                Runtime runtime = Runtime.getRuntime();
                
                for (int i = 0; i < 60; i++) {
                    long usedMemory = (runtime.totalMemory() - runtime.freeMemory()) / 1024 / 1024;
                    long totalMemory = runtime.totalMemory() / 1024 / 1024;
                    int processors = runtime.availableProcessors();
                    
                    emitter.send(SseEmitter.event()
                        .name("status")
                        .data(Map.of(
                            "usedMemoryMB", usedMemory,
                            "totalMemoryMB", totalMemory,
                            "processors", processors,
                            "timestamp", Instant.now()
                        ))
                    );
                    
                    Thread.sleep(1000);
                }
                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });
        
        return emitter;
    }
    
    /**
     * 发送通知给指定用户
     */
    @PostMapping("/send-notification")
    public Map<String, Object> sendNotification(
            @RequestParam String userId,
            @RequestBody Notification notification) {
        
        SseEmitter emitter = emitters.get(userId);
        
        if (emitter == null) {
            return Map.of("success", false, "message", "User not connected");
        }
        
        try {
            emitter.send(SseEmitter.event()
                .name("notification")
                .data(notification)
            );
            return Map.of("success", true);
        } catch (IOException e) {
            log.error("Failed to send notification", e);
            emitters.remove(userId);
            return Map.of("success", false, "message", e.getMessage());
        }
    }
    
    /**
     * 广播通知给所有连接的用户
     */
    @PostMapping("/broadcast")
    public Map<String, Object> broadcast(@RequestBody Notification notification) {
        int successCount = 0;
        int failCount = 0;
        
        for (Map.Entry<String, SseEmitter> entry : emitters.entrySet()) {
            try {
                entry.getValue().send(SseEmitter.event()
                    .name("notification")
                    .data(notification)
                );
                successCount++;
            } catch (IOException e) {
                log.error("Failed to send to user: {}", entry.getKey(), e);
                emitters.remove(entry.getKey());
                failCount++;
            }
        }
        
        return Map.of(
            "totalConnections", emitters.size(),
            "successCount", successCount,
            "failCount", failCount
        );
    }
    
    /**
     * 获取当前连接数
     */
    @GetMapping("/connections")
    public Map<String, Object> getConnections() {
        return Map.of(
            "activeConnections", emitters.size(),
            "users", emitters.keySet()
        );
    }
}

<!DOCTYPE html>
<html>
<head>
    <title>SSE Demo - Spring Boot 4</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 1200px;
            margin: 50px auto;
            padding: 20px;
        }
        .container {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 20px;
        }
        .panel {
            border: 1px solid #ccc;
            padding: 15px;
            border-radius: 5px;
        }
        .events {
            height: 300px;
            overflow-y: scroll;
            border: 1px solid #eee;
            padding: 10px;
            margin-top: 10px;
        }
        .event {
            margin: 5px 0;
            padding: 5px;
            background: #f5f5f5;
        }
        button {
            padding: 10px 20px;
            margin: 5px;
        }
    </style>
</head>
<body>
    <h1>Server-Sent Events Demo</h1>
    
    <div class="container">
        <div class="panel">
            <h2>Notifications</h2>
            <input type="text" id="userId" placeholder="User ID" value="user123" />
            <button onclick="connectNotifications()">Connect</button>
            <button onclick="disconnectNotifications()">Disconnect</button>
            <div id="notifications" class="events"></div>
        </div>
        
        <div class="panel">
            <h2>Stock Prices</h2>
            <input type="text" id="symbol" placeholder="Symbol" value="AAPL" />
            <button onclick="connectStockPrices()">Connect</button>
            <button onclick="disconnectStockPrices()">Disconnect</button>
            <div id="stockPrices" class="events"></div>
        </div>
        
        <div class="panel">
            <h2>Server Status</h2>
            <button onclick="connectServerStatus()">Connect</button>
            <button onclick="disconnectServerStatus()">Disconnect</button>
            <div id="serverStatus" class="events"></div>
        </div>
    </div>
    
    <script>
        let notificationSource = null;
        let stockPriceSource = null;
        let serverStatusSource = null;
        
        function connectNotifications() {
            const userId = document.getElementById('userId').value;
            const url = `http://localhost:8080/api/sse/notifications?userId=${userId}`;
            
            notificationSource = new EventSource(url);
            
            notificationSource.addEventListener('connected', (e) => {
                addEvent('notifications', 'Connected: ' + e.data);
            });
            
            notificationSource.addEventListener('notification', (e) => {
                const data = JSON.parse(e.data);
                addEvent('notifications', `[${data.type}] ${data.title}: ${data.message}`);
            });
            
            notificationSource.onerror = () => {
                addEvent('notifications', 'Connection error');
            };
        }
        
        function disconnectNotifications() {
            if (notificationSource) {
                notificationSource.close();
                notificationSource = null;
                addEvent('notifications', 'Disconnected');
            }
        }
        
        function connectStockPrices() {
            const symbol = document.getElementById('symbol').value;
            const url = `http://localhost:8080/api/sse/stock-prices?symbol=${symbol}`;
            
            stockPriceSource = new EventSource(url);
            
            stockPriceSource.addEventListener('price-update', (e) => {
                const data = JSON.parse(e.data);
                addEvent('stockPrices', `${data.symbol}: $${data.price}`);
            });
            
            stockPriceSource.onerror = () => {
                addEvent('stockPrices', 'Connection error');
            };
        }
        
        function disconnectStockPrices() {
            if (stockPriceSource) {
                stockPriceSource.close();
                stockPriceSource = null;
                addEvent('stockPrices', 'Disconnected');
            }
        }
        
        function connectServerStatus() {
            const url = 'http://localhost:8080/api/sse/server-status';
            
            serverStatusSource = new EventSource(url);
            
            serverStatusSource.addEventListener('status', (e) => {
                const data = JSON.parse(e.data);
                addEvent('serverStatus', 
                    `Memory: ${data.usedMemoryMB}/${data.totalMemoryMB}MB, CPUs: ${data.processors}`);
            });
            
            serverStatusSource.onerror = () => {
                addEvent('serverStatus', 'Connection error');
            };
        }
        
        function disconnectServerStatus() {
            if (serverStatusSource) {
                serverStatusSource.close();
                serverStatusSource = null;
                addEvent('serverStatus', 'Disconnected');
            }
        }
        
        function addEvent(containerId, text) {
            const container = document.getElementById(containerId);
            const eventDiv = document.createElement('div');
            eventDiv.className = 'event';
            eventDiv.textContent = new Date().toLocaleTimeString() + ' - ' + text;
            container.appendChild(eventDiv);
            container.scrollTop = container.scrollHeight;
            
            // 限制显示的事件数量
            while (container.children.length > 50) {
                container.removeChild(container.firstChild);
            }
        }
    </script>
</body>
</html>