WebSocket Support

/**
 * WebSocket connection abstraction
 */
public interface WebSocket {
    /**
     * Write string data to WebSocket
     * @param data string data to send
     * @return this WebSocket for chaining
     */
    public WebSocket write(String data);
    
    /**
     * Write byte array data to WebSocket
     * @param data byte array to send
     * @return this WebSocket for chaining
     */
    public WebSocket write(byte[] data);
    
    /**
     * Write generic object data (will be serialized)
     * @param data object to send
     * @return this WebSocket for chaining
     */
    public WebSocket write(Object data);
    
    /**
     * Close the WebSocket connection
     */
    public void close();
    
    /**
     * Close with specific code and reason
     * @param closeCode WebSocket close code
     * @param message close reason message
     */
    public void close(int closeCode, String message);
    
    /**
     * Check if WebSocket is open and ready
     * @return true if connection is open
     */
    public boolean isOpen();
    
    /**
     * Get the associated AtmosphereResource
     * @return AtmosphereResource instance
     */
    public AtmosphereResource resource();
}

public class ChatWebSocketHandler implements WebSocketHandler {
    @Override
    public void onOpen(WebSocket webSocket) {
        System.out.println("WebSocket opened");
        webSocket.write("Welcome to chat!");
    }
    
    @Override
    public void onTextMessage(WebSocket webSocket, String message) {
        // Echo message back
        webSocket.write("Echo: " + message);
        
        // Broadcast to all connected clients
        webSocket.resource().getBroadcaster().broadcast(message);
    }
    
    @Override
    public void onClose(WebSocket webSocket) {
        System.out.println("WebSocket closed");
    }
}

/**
 * Handle WebSocket-specific events and messages
 */
public interface WebSocketHandler {
    /**
     * Called when WebSocket connection is opened
     * @param webSocket the opened WebSocket
     */
    public void onOpen(WebSocket webSocket);
    
    /**
     * Called when WebSocket connection is closed
     * @param webSocket the closed WebSocket
     */
    public void onClose(WebSocket webSocket);
    
    /**
     * Called when WebSocket connection encounters an error
     * @param webSocket the WebSocket with error
     * @param t the error that occurred
     */
    public void onError(WebSocket webSocket, Throwable t);
    
    /**
     * Called when text message is received
     * @param webSocket the receiving WebSocket
     * @param message the text message received
     */
    public void onTextMessage(WebSocket webSocket, String message);
    
    /**
     * Called when binary message is received
     * @param webSocket the receiving WebSocket
     * @param message the binary message as byte array
     */
    public void onByteMessage(WebSocket webSocket, byte[] message);
}

@WebSocketHandlerService(path = "/websocket/data")
public class DataWebSocketHandler implements WebSocketHandler {
    
    @Override
    public void onOpen(WebSocket webSocket) {
        // Send initial data when client connects
        webSocket.write("{\"type\":\"welcome\",\"message\":\"Connected\"}");
        
        // Add to broadcaster for group messages
        webSocket.resource().getBroadcaster()
                .addAtmosphereResource(webSocket.resource());
    }
    
    @Override
    public void onTextMessage(WebSocket webSocket, String message) {
        try {
            // Parse JSON message
            JsonObject json = JsonParser.parseString(message).getAsJsonObject();
            String type = json.get("type").getAsString();
            
            switch (type) {
                case "ping":
                    webSocket.write("{\"type\":\"pong\"}");
                    break;
                case "broadcast":
                    String msg = json.get("message").getAsString();
                    webSocket.resource().getBroadcaster().broadcast(msg);
                    break;
                default:
                    webSocket.write("{\"type\":\"error\",\"message\":\"Unknown type\"}");
            }
        } catch (Exception e) {
            webSocket.write("{\"type\":\"error\",\"message\":\"Invalid JSON\"}");
        }
    }
    
    @Override
    public void onByteMessage(WebSocket webSocket, byte[] message) {
        // Handle binary data (e.g., file uploads, images)
        System.out.println("Received " + message.length + " bytes");
        
        // Echo binary data back
        webSocket.write(message);
    }
    
    @Override
    public void onClose(WebSocket webSocket) {
        // Cleanup when client disconnects
        webSocket.resource().getBroadcaster()
                .removeAtmosphereResource(webSocket.resource());
        System.out.println("WebSocket closed for: " + webSocket.resource().uuid());
    }
    
    @Override
    public void onError(WebSocket webSocket, Throwable t) {
        System.err.println("WebSocket error: " + t.getMessage());
        t.printStackTrace();
    }
}

/**
 * Process WebSocket messages and manage protocol lifecycle
 */
public interface WebSocketProcessor {
    /**
     * Open WebSocket connection
     * @param webSocket the WebSocket to open
     * @return WebSocket instance
     */
    public WebSocket open(WebSocket webSocket);
    
    /**
     * Invoke WebSocket protocol handling
     * @param webSocket the WebSocket
     * @param webSocketMessage message to process
     */
    public void invokeWebSocketProtocol(WebSocket webSocket, String webSocketMessage);
    
    /**
     * Close WebSocket connection
     * @param webSocket the WebSocket to close
     * @param closeCode close status code
     */
    public void close(WebSocket webSocket, int closeCode);
    
    /**
     * Destroy processor and cleanup resources
     */
    public void destroy();
}

/**
 * Default WebSocket message processor
 */
public class DefaultWebSocketProcessor implements WebSocketProcessor {
    /**
     * Create processor with AtmosphereFramework
     * @param framework AtmosphereFramework instance
     */
    public DefaultWebSocketProcessor(AtmosphereFramework framework);
    
    /**
     * Register WebSocketHandler for path pattern
     * @param path URL path pattern
     * @param handler WebSocketHandler instance
     * @return this processor
     */
    public DefaultWebSocketProcessor registerWebSocketHandler(String path, WebSocketHandler handler);
}

/**
 * Define WebSocket message protocols
 */
public interface WebSocketProtocol {
    /**
     * Handle incoming WebSocket message according to protocol
     * @param webSocket the WebSocket connection
     * @param message the incoming message
     * @return list of processed messages for broadcast
     */
    public List<AtmosphereRequest> onMessage(WebSocket webSocket, String message);
    
    /**
     * Configure the protocol
     * @param config AtmosphereConfig for setup
     */
    public void configure(AtmosphereConfig config);
    
    /**
     * Get supported protocols
     * @return list of supported protocol names
     */
    public List<String> supportedProtocols();
    
    /**
     * Inspect message and determine if protocol applies
     * @param message message to inspect
     * @return true if this protocol handles the message
     */
    public boolean inspectResponse();
}

/**
 * Simple HTTP-over-WebSocket protocol
 */
public class SimpleHttpProtocol implements WebSocketProtocol {
    /**
     * Create protocol with delimiter
     * @param delimiter message delimiter character
     */
    public SimpleHttpProtocol(String delimiter);
}

/**
 * Streaming HTTP protocol for WebSocket
 */
public class StreamingHttpProtocol implements WebSocketProtocol {
    /**
     * Create streaming protocol with delimiter
     * @param delimiter message delimiter
     */
    public StreamingHttpProtocol(String delimiter);
}

// Register custom WebSocket protocol
@WebSocketProtocolService
public class JsonWebSocketProtocol implements WebSocketProtocol {
    
    @Override
    public List<AtmosphereRequest> onMessage(WebSocket webSocket, String message) {
        List<AtmosphereRequest> requests = new ArrayList<>();
        
        try {
            // Parse JSON message
            JsonObject json = JsonParser.parseString(message).getAsJsonObject();
            
            // Create AtmosphereRequest from JSON
            AtmosphereRequest.Builder builder = new AtmosphereRequest.Builder();
            builder.pathInfo(json.get("path").getAsString());
            builder.method(json.get("method").getAsString());
            
            if (json.has("body")) {
                builder.body(json.get("body").getAsString());
            }
            
            requests.add(builder.build());
            
        } catch (Exception e) {
            // Send error response
            webSocket.write("{\"error\":\"Invalid JSON protocol message\"}");
        }
        
        return requests;
    }
    
    @Override
    public void configure(AtmosphereConfig config) {
        // Protocol configuration
    }
    
    @Override
    public List<String> supportedProtocols() {
        return Arrays.asList("json-protocol");
    }
}

/**
 * Factory for creating WebSocket instances
 */
public interface WebSocketFactory {
    /**
     * Create WebSocket instance
     * @param request AtmosphereRequest
     * @param response AtmosphereResponse
     * @return WebSocket instance
     */
    public WebSocket create(AtmosphereRequest request, AtmosphereResponse response);
    
    /**
     * Find existing WebSocket by AtmosphereResource
     * @param resource AtmosphereResource
     * @return WebSocket instance or null
     */
    public WebSocket find(AtmosphereResource resource);
}

/**
 * Default WebSocket factory implementation
 */
public class DefaultWebSocketFactory implements WebSocketFactory {
    /**
     * Create factory with AtmosphereConfig
     * @param config AtmosphereConfig instance
     */
    public DefaultWebSocketFactory(AtmosphereConfig config);
}

/**
 * WebSocket event listener for monitoring connections
 */
public interface WebSocketEventListener {
    /**
     * Called when any WebSocket opens
     * @param event WebSocket event
     */
    public void onWebSocketConnect(WebSocketEvent event);
    
    /**
     * Called when any WebSocket closes
     * @param event WebSocket event
     */
    public void onWebSocketClose(WebSocketEvent event);
    
    /**
     * Called on WebSocket error
     * @param event WebSocket event with error
     */
    public void onWebSocketError(WebSocketEvent event);
}

/**
 * WebSocket event containing connection and context information
 */
public class WebSocketEvent {
    /**
     * Get the WebSocket connection
     * @return WebSocket instance
     */
    public WebSocket webSocket();
    
    /**
     * Get any associated error
     * @return Throwable or null
     */
    public Throwable throwable();
    
    /**
     * Get event type
     * @return WebSocketEventType enum
     */
    public WebSocketEventType type();
}

/**
 * Types of WebSocket events
 */
public enum WebSocketEventType {
    CONNECT,
    MESSAGE, 
    CLOSE,
    ERROR
}

/**
 * Asynchronous I/O writer for non-blocking writes
 */
public interface AsyncIOWriter {
    /**
     * Write data asynchronously
     * @param resource AtmosphereResource target
     * @param data data to write
     * @return Future for async write operation
     */
    public Future<String> write(AtmosphereResource resource, String data);
    
    /**
     * Write byte data asynchronously
     * @param resource AtmosphereResource target
     * @param data byte array to write
     * @return Future for async write operation
     */
    public Future<byte[]> write(AtmosphereResource resource, byte[] data);
    
    /**
     * Write generic object asynchronously
     * @param resource AtmosphereResource target
     * @param data object to write (will be serialized)
     * @return Future for async write operation
     */
    public Future<Object> write(AtmosphereResource resource, Object data);
    
    /**
     * Close the writer and cleanup resources
     * @param resource AtmosphereResource
     */
    public void close(AtmosphereResource resource);
    
    /**
     * Flush pending writes
     * @param resource AtmosphereResource
     */
    public void flush(AtmosphereResource resource);
}

/**
 * Default AsyncIOWriter implementation
 */
public class DefaultAsyncIOWriter implements AsyncIOWriter {
    /**
     * Create writer with AtmosphereResponse
     * @param response AtmosphereResponse for writing
     */
    public DefaultAsyncIOWriter(AtmosphereResponse response);
    
    /**
     * Set write timeout
     * @param timeout write timeout in milliseconds
     * @return this writer
     */
    public DefaultAsyncIOWriter setTimeout(long timeout);
}

@WebSocketHandlerService(path = "/async-websocket")
public class AsyncWebSocketHandler implements WebSocketHandler {
    
    @Override
    public void onOpen(WebSocket webSocket) {
        AtmosphereResource resource = webSocket.resource();
        AsyncIOWriter writer = resource.getResponse().getAsyncIOWriter();
        
        // Asynchronous welcome message
        Future<String> writeResult = writer.write(resource, "Welcome! Connection established.");
        
        writeResult.thenAccept(result -> {
            System.out.println("Welcome message sent successfully");
        }).exceptionally(throwable -> {
            System.err.println("Failed to send welcome message: " + throwable.getMessage());
            return null;
        });
    }
    
    @Override
    public void onTextMessage(WebSocket webSocket, String message) {
        AtmosphereResource resource = webSocket.resource();
        AsyncIOWriter writer = resource.getResponse().getAsyncIOWriter();
        
        // Process message asynchronously
        CompletableFuture.supplyAsync(() -> processMessage(message))
            .thenCompose(processed -> writer.write(resource, processed))
            .thenAccept(result -> System.out.println("Response sent"))
            .exceptionally(throwable -> {
                System.err.println("Error processing message: " + throwable.getMessage());
                return null;
            });
    }
    
    private String processMessage(String message) {
        // Simulate processing
        try {
            Thread.sleep(100); // Simulate work
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        return "Processed: " + message;
    }
}

@WebSocketHandlerService(path = "/websocket/chat/{room}")
public class ChatWebSocketHandler implements WebSocketHandler {
    private final Map<String, Set<WebSocket>> roomConnections = new ConcurrentHashMap<>();
    
    @Override
    public void onOpen(WebSocket webSocket) {
        // Extract room from path
        String room = extractRoomFromPath(webSocket.resource().getRequest().getPathInfo());
        
        // Add to room connections
        roomConnections.computeIfAbsent(room, k -> ConcurrentHashMap.newKeySet())
                      .add(webSocket);
        
        // Send welcome message
        webSocket.write(createMessage("system", "Welcome to room " + room));
        
        // Notify room of new user
        broadcastToRoom(room, createMessage("system", "User joined the room"), webSocket);
        
        System.out.println("WebSocket opened for room: " + room);
    }
    
    @Override
    public void onTextMessage(WebSocket webSocket, String message) {
        String room = extractRoomFromPath(webSocket.resource().getRequest().getPathInfo());
        
        try {
            // Parse incoming message
            JsonObject json = JsonParser.parseString(message).getAsJsonObject();
            String type = json.get("type").getAsString();
            
            switch (type) {
                case "chat":
                    String text = json.get("text").getAsString();
                    String user = json.get("user").getAsString();
                    
                    // Broadcast chat message to room
                    String chatMessage = createMessage("chat", text, user);
                    broadcastToRoom(room, chatMessage, null); // Include sender
                    break;
                    
                case "typing":
                    String typingUser = json.get("user").getAsString();
                    String typingMessage = createMessage("typing", typingUser + " is typing...");
                    broadcastToRoom(room, typingMessage, webSocket); // Exclude sender
                    break;
                    
                default:
                    webSocket.write(createMessage("error", "Unknown message type"));
            }
        } catch (Exception e) {
            webSocket.write(createMessage("error", "Invalid message format"));
        }
    }
    
    @Override
    public void onClose(WebSocket webSocket) {
        String room = extractRoomFromPath(webSocket.resource().getRequest().getPathInfo());
        
        // Remove from room connections
        Set<WebSocket> connections = roomConnections.get(room);
        if (connections != null) {
            connections.remove(webSocket);
            if (connections.isEmpty()) {
                roomConnections.remove(room);
            }
        }
        
        // Notify room of user leaving
        broadcastToRoom(room, createMessage("system", "User left the room"), null);
        
        System.out.println("WebSocket closed for room: " + room);
    }
    
    @Override
    public void onError(WebSocket webSocket, Throwable t) {
        System.err.println("WebSocket error: " + t.getMessage());
        t.printStackTrace();
        
        // Send error message to client
        webSocket.write(createMessage("error", "Connection error occurred"));
    }
    
    private void broadcastToRoom(String room, String message, WebSocket exclude) {
        Set<WebSocket> connections = roomConnections.get(room);
        if (connections != null) {
            connections.stream()
                     .filter(ws -> ws != exclude && ws.isOpen())
                     .forEach(ws -> ws.write(message));
        }
    }
    
    private String createMessage(String type, String content) {
        return String.format("{\"type\":\"%s\",\"content\":\"%s\",\"timestamp\":%d}", 
                           type, content, System.currentTimeMillis());
    }
    
    private String createMessage(String type, String content, String user) {
        return String.format("{\"type\":\"%s\",\"content\":\"%s\",\"user\":\"%s\",\"timestamp\":%d}", 
                           type, content, user, System.currentTimeMillis());
    }
    
    private String extractRoomFromPath(String pathInfo) {
        // Extract room from path like "/websocket/chat/room123"
        String[] parts = pathInfo.split("/");
        return parts.length > 3 ? parts[3] : "default";
    }
}