package ru.dragonestia.kubik.discovery.ws;

import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import ru.dragonestia.kubik.discovery.component.UpdateStateNotifier;
import tools.jackson.databind.ObjectMapper;

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

@Log4j2
@RequiredArgsConstructor
public class WebSocketHandler extends TextWebSocketHandler {

    private final UpdateStateNotifier updateStateNotifier;
    private final ObjectMapper objectMapper;

    private final Map<String, UpdateStateNotifier.Subscriber> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(@NonNull WebSocketSession session) {
        log.info("Opened new ws connection form {}", session.getRemoteAddress());

        var subscriber = updateStateNotifier.subscribe(response -> {
            try {
                send(session, objectMapper.writeValueAsString(response));
            } catch (Exception ex) {
                log.error(ex.getMessage(), ex);
            }
        });
        sessions.put(session.getId(), subscriber);
    }

    @Override
    public void afterConnectionClosed(@NonNull WebSocketSession session, @NonNull CloseStatus status) {
        log.info("Closed ws connection form {}", session.getRemoteAddress());

        var subscriber = sessions.remove(session.getId());
        if (subscriber != null) subscriber.dispose();
    }

    private void send(WebSocketSession session, String message) {
        log.debug("Sending WS message to {}: {}", session.getRemoteAddress(), message);

        try {
            session.sendMessage(new TextMessage(message));
        } catch (Exception ex) {
            log.error(ex.getMessage(), ex);
        }
    }
}
