package fr.redsavant.ws;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.WebSocket;
import java.nio.ByteBuffer;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.logging.Logger;

public class WebSocketClient implements WebSocket.Listener {

    private static final Logger LOGGER = Logger.getLogger(WebSocketClient.class.getName());

    private final URI uri;
    private final Consumer<String> onMessage;
    private WebSocket socket;
    private boolean connected = false;
    private final ScheduledExecutorService reconnectScheduler = Executors.newSingleThreadScheduledExecutor();

    public WebSocketClient(URI uri, Consumer<String> onMessage) {
        this.uri = uri;
        this.onMessage = onMessage;
    }

    public void connect() {
        HttpClient.newHttpClient()
                .newWebSocketBuilder()
                .buildAsync(uri, this)
                .thenAccept(ws -> {
                    this.socket = ws;
                    this.connected = true;
                    LOGGER.info("WebSocket connected : " + uri);
                })
                .exceptionally(ex -> {
                    LOGGER.warning("Failed to connect to WebSocket, retry in 5s : " + ex.getMessage());
                    scheduleReconnect();
                    return null;
                });
    }

    public void send(String message) {
        if (socket != null && connected) {
            socket.sendText(message, true);
        }
    }

    public void close() {
        reconnectScheduler.shutdownNow();
        if (socket != null) {
            socket.sendClose(WebSocket.NORMAL_CLOSURE, "shutdown");
        }
    }

    public boolean isConnected() {
        return connected;
    }

    private void scheduleReconnect() {
        reconnectScheduler.schedule(this::connect, 5, TimeUnit.SECONDS);
    }

    @Override
    public void onOpen(WebSocket webSocket) {
        connected = true;
        LOGGER.info("WebSocket opened");
        webSocket.request(1);
    }

    @Override
    public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) {
        onMessage.accept(data.toString());
        webSocket.request(1);
        return null;
    }

    @Override
    public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
        connected = false;
        LOGGER.warning("WebSocket closed (" + statusCode + "), re-connect in 5s...");
        scheduleReconnect();
        return null;
    }

    @Override
    public void onError(WebSocket webSocket, Throwable error) {
        connected = false;
        LOGGER.warning("WebSocket error : " + error.getMessage());
        scheduleReconnect();
    }
}