/*
 * Javalin - https://javalin.io
 * Copyright 2017 David Åse
 * Licensed under Apache 2.0: https://github.com/tipsy/javalin/blob/master/LICENSE
 */

package io.javalin.embeddedserver.jetty.websocket;

import io.javalin.embeddedserver.jetty.websocket.interfaces.CloseHandler;
import io.javalin.embeddedserver.jetty.websocket.interfaces.ConnectHandler;
import io.javalin.embeddedserver.jetty.websocket.interfaces.ErrorHandler;
import io.javalin.embeddedserver.jetty.websocket.interfaces.MessageHandler;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.stream.Collectors;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.annotations.OnWebSocketClose;
import org.eclipse.jetty.websocket.api.annotations.OnWebSocketConnect;
import org.eclipse.jetty.websocket.api.annotations.OnWebSocketError;
import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage;
import org.eclipse.jetty.websocket.api.annotations.WebSocket;
import org.jetbrains.annotations.NotNull;

@WebSocket
public class WebSocketHandler {

    private final ConcurrentMap<Session, String> sessions = new ConcurrentHashMap<>();

    private ConnectHandler connectHandler = null;
    private MessageHandler messageHandler = null;
    private CloseHandler closeHandler = null;
    private ErrorHandler errorHandler = null;

    public void onConnect(@NotNull ConnectHandler connectHandler) {
        this.connectHandler = connectHandler;
    }

    public void onMessage(@NotNull MessageHandler messageHandler) {
        this.messageHandler = messageHandler;
    }

    public void onClose(@NotNull CloseHandler closeHandler) {
        this.closeHandler = closeHandler;
    }

    public void onError(@NotNull ErrorHandler errorHandler) {
        this.errorHandler = errorHandler;
    }

    /**
     * Returns all open sessions at the time of invocation. If you intend to use this method to build a client broadcast
     * capability then you should handle connection errors due to difference in time between retrieval and
     * {@link WsSession#send(String)}.
     *
     * @return a collection of {@link WsSession}
     */
    public Set<WsSession> getSessions() {
        return sessions.keySet()
            .stream()
            .filter(Session::isOpen)
            .map(this::createWsSession)
            .collect(Collectors.toSet());
    }

    // Jetty annotations

    @OnWebSocketConnect
    public void _internalOnConnectProxy(Session session) throws Exception {
        registerSession(session);
        if (connectHandler != null) {
            connectHandler.handle(createWsSession(session));
        }
    }

    @OnWebSocketMessage
    public void _internalOnMessageProxy(Session session, String message) throws Exception {
        registerSession(session);
        if (messageHandler != null) {
            messageHandler.handle(createWsSession(session), message);
        }
    }

    @OnWebSocketClose
    public void _internalOnCloseProxy(Session session, int statusCode, String reason) throws Exception {
        if (closeHandler != null) {
            closeHandler.handle(createWsSession(session), statusCode, reason);
        }
        unregisterSession(session);
    }

    @OnWebSocketError
    public void _internalOnErrorProxy(Session session, Throwable throwable) throws Exception {
        registerSession(session);
        if (errorHandler != null) {
            errorHandler.handle(createWsSession(session), throwable);
        }
    }

    private void registerSession(Session session) {
        sessions.putIfAbsent(session, UUID.randomUUID().toString());
    }

    private void unregisterSession(Session session) {
        sessions.remove(session);
    }

    private WsSession createWsSession(Session session) {
        return new WsSession(sessions.get(session), session);
    }

}

@@ -2385,24 +2385,31 @@
 p(this::
-c
 re
-ateWs
+gisterAndWrap
 Session)
@@ -2583,32 +2583,54 @@
 eption %7B%0A       
+ WsSession wsSession =
  registerSession
@@ -2614,32 +2614,39 @@
 ssion = register
+AndWrap
 Session(session)
@@ -2715,39 +2715,33 @@
 dler.handle(
-createW
+w
 sSession
 (session));%0A
@@ -2728,25 +2728,16 @@
 sSession
-(session)
 );%0A     
@@ -2863,32 +2863,54 @@
 eption %7B%0A       
+ WsSession wsSession =
  registerSession
@@ -2894,32 +2894,39 @@
 ssion = register
+AndWrap
 Session(session)
@@ -2995,39 +2995,33 @@
 dler.handle(
-createW
+w
 sSession
 (session), m
@@ -3008,25 +3008,16 @@
 sSession
-(session)
 , messag
@@ -3156,32 +3156,95 @@
 ows Exception %7B%0A
+        WsSession wsSession = registerAndWrapSession(session);%0A
         if (clos
@@ -3291,39 +3291,33 @@
 dler.handle(
-createW
+w
 sSession
 (session), s
@@ -3304,25 +3304,16 @@
 sSession
-(session)
 , status
@@ -3345,33 +3345,31 @@
         
-unregisterSession
+sessions.remove
 (session
@@ -3502,16 +3502,38 @@
 %0A       
+ WsSession wsSession =
  registe
@@ -3525,32 +3525,39 @@
 ssion = register
+AndWrap
 Session(session)
@@ -3626,31 +3626,25 @@
 .handle(
-createW
+w
 sSession
 (session
@@ -3635,25 +3635,16 @@
 sSession
-(session)
 , throwa
@@ -3678,20 +3678,25 @@
 private 
-void
+WsSession
  registe
@@ -3688,32 +3688,39 @@
 Session register
+AndWrap
 Session(Session 
@@ -3803,167 +3803,8 @@
 ));%0A
-    %7D%0A%0A    private void unregisterSession(Session session) %7B%0A        sessions.remove(session);%0A    %7D%0A%0A    private WsSession createWsSession(Session session) %7B%0A
     

Simplify websocket session tracking