path: root/src/main/java/dev/plutorocks/BridgeClient.java

package dev.plutorocks;

import com.google.gson.Gson;
import com.google.gson.JsonObject;

import net.minecraft.client.MinecraftClient;
import net.minecraft.text.MutableText;
import net.minecraft.text.Text;
import net.minecraft.util.Formatting;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.WebSocket;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicBoolean;

public class BridgeClient {
    private static final Logger LOGGER = LoggerFactory.getLogger("PlutoBridgeClient");
    private static final Gson GSON = new Gson();

    private final String host;
    private final String token;
    private volatile String nick;
    private final String serverId = "default";

    private final AtomicBoolean running = new AtomicBoolean(false);
    private volatile boolean connected = false;

    private WebSocket webSocket;

    public BridgeClient(String host, String token, String nick) {
        this.host = host;
        this.token = token;
        this.nick = (nick == null || nick.isEmpty()) ? "Player" : nick;
    }

    public boolean isConnected() {
        return connected;
    }

    public void connect() {
        if (running.getAndSet(true)) {
            return; // already running
        }

        try {
            String uriStr;

            // if user already gave us a full WebSocket URL, trust it.
            if (host.startsWith("ws://") || host.startsWith("wss://")) {
                uriStr = host;
            } else if (host.startsWith("http://") || host.startsWith("https://")) {
                // convert http(s) -> ws(s)
                boolean secure = host.startsWith("https://");
                String rest = host.substring(secure ? "https://".length() : "http://".length());
                uriStr = (secure ? "wss://" : "ws://") + rest;

                // if there's no path after the host, add /ws
                if (!rest.contains("/")) {
                    uriStr = uriStr + "/ws";
                }
            } else {
                // bare hostname: assume HTTPS/WSS + /ws
                // this matches typical reverse-proxy setup (TLS on 443).
                uriStr = "wss://" + host;
                if (!host.contains("/")) {
                    uriStr = uriStr + "/ws";
                }
            }

            LOGGER.info("Connecting to bridge at {}", uriStr);

            HttpClient client = HttpClient.newHttpClient();
            WebSocket.Builder builder = client.newWebSocketBuilder();
            CompletableFuture<WebSocket> future = builder.buildAsync(URI.create(uriStr), new BridgeListener());
            future.whenComplete((ws, throwable) -> {
                if (throwable != null) {
                    LOGGER.error("Failed to connect to bridge", throwable);
                    running.set(false);
                    String msg = throwable.getMessage();
                    if (msg == null) msg = throwable.getClass().getSimpleName();
                    sendClientChat("Failed to connect to bridge: " + msg);
                    return;
                }
                this.webSocket = ws;
                sendAuth();
            });
        } catch (Exception e) {
            LOGGER.error("Error starting WebSocket connection", e);
            running.set(false);
            sendClientChat("Bridge connection error: " + e.getMessage());
        }
    }

    public void disconnect() {
        running.set(false);
        connected = false;
        if (webSocket != null) {
            try {
                webSocket.sendClose(WebSocket.NORMAL_CLOSURE, "client disconnect");
            } catch (Exception ignored) {}
        }
    }

    public void sendChatMessage(String message) {
        if (!connected || webSocket == null) {
            sendClientChat("Bridge is not connected.");
            return;
        }

        JsonObject obj = new JsonObject();
        obj.addProperty("type", "chat");
        obj.addProperty("message", message);

        String json = GSON.toJson(obj);
        webSocket.sendText(json, true);
    }

    public void updateNick(String newNick) {
        this.nick = newNick;
        if (!connected || webSocket == null) {
            return;
        }

        JsonObject obj = new JsonObject();
        obj.addProperty("type", "nick");
        obj.addProperty("nick", newNick);

        String json = GSON.toJson(obj);
        webSocket.sendText(json, true);
    }

    private void sendAuth() {
        if (webSocket == null) {
            return;
        }

        JsonObject auth = new JsonObject();
        auth.addProperty("type", "auth");
        auth.addProperty("token", token);
        auth.addProperty("nick", nick);
        auth.addProperty("serverId", serverId);
        // playerUuid optional

        String json = GSON.toJson(auth);
        webSocket.sendText(json, true);
    }

    private void sendClientChat(String message) {
        MinecraftClient client = MinecraftClient.getInstance();
        if (client == null || client.inGameHud == null) {
            return;
        }
        MutableText prefix = Text.literal("[Bridge] ").formatted(Formatting.AQUA);
        MutableText body = Text.literal(message).formatted(Formatting.GRAY);
        client.execute(() -> client.inGameHud.getChatHud().addMessage(prefix.append(body)));
    }

    private void handleIncomingChat(JsonObject obj) {
        String fromNick = obj.has("fromNick") ? obj.get("fromNick").getAsString() : "?";
        String msg = obj.has("message") ? obj.get("message").getAsString() : "";

        MinecraftClient client = MinecraftClient.getInstance();
        if (client == null || client.inGameHud == null) {
            return;
        }

        MutableText line = Text.literal("[Bridge] ")
                .formatted(Formatting.AQUA)
                .append(Text.literal("<" + fromNick + "> ").formatted(Formatting.GRAY))
                .append(Text.literal(msg).formatted(Formatting.WHITE));

        client.execute(() -> client.inGameHud.getChatHud().addMessage(line));
    }

    private class BridgeListener implements WebSocket.Listener {

        private final StringBuilder textBuffer = new StringBuilder();

        @Override
        public void onOpen(WebSocket webSocket) {
            LOGGER.info("Bridge WebSocket opened");
            connected = true;
            sendClientChat("Connected to bridge.");
            WebSocket.Listener.super.onOpen(webSocket);
        }

        @Override
        public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) {
            textBuffer.append(data);
            if (last) {
                String full = textBuffer.toString();
                textBuffer.setLength(0);

                try {
                    JsonObject obj = GSON.fromJson(full, JsonObject.class);
                    if (obj != null && obj.has("type")) {
                        String type = obj.get("type").getAsString();
                        if ("chat".equals(type)) {
                            handleIncomingChat(obj);
                        }
                    }
                } catch (Exception e) {
                    LOGGER.warn("Failed to parse bridge message: {}", full, e);
                }
            }
            return WebSocket.Listener.super.onText(webSocket, data, last);
        }

        @Override
        public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
            LOGGER.info("Bridge WebSocket closed: {} {}", statusCode, reason);
            connected = false;
            running.set(false);

            String reasonText = (reason == null ? "" : reason).toLowerCase();

            // 1008 = Policy Violation -> we use this for "token revoked"
            if (statusCode == 1008 && reasonText.contains("token revoked")) {
                sendClientChat("Disconnected from bridge: token revoked by server.");
            } else {
                sendClientChat("Disconnected from bridge (" + statusCode + "): " + reason);
            }

            return CompletableFuture.completedFuture(null);
        }

        @Override
        public void onError(WebSocket webSocket, Throwable error) {
            LOGGER.error("Bridge WebSocket error", error);
            connected = false;
            running.set(false);
            sendClientChat("Bridge connection error: " + error.getMessage());
        }

        @Override
        public CompletionStage<?> onBinary(WebSocket webSocket, java.nio.ByteBuffer data, boolean last) {
            return WebSocket.Listener.super.onBinary(webSocket, data, last);
        }

        @Override
        public CompletionStage<?> onPong(WebSocket webSocket, java.nio.ByteBuffer message) {
            return WebSocket.Listener.super.onPong(webSocket, message);
        }

        @Override
        public CompletionStage<?> onPing(WebSocket webSocket, java.nio.ByteBuffer message) {
            webSocket.sendPong(java.nio.ByteBuffer.wrap("pong".getBytes(StandardCharsets.UTF_8)));
            return WebSocket.Listener.super.onPing(webSocket, message);
        }
    }
}