import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UncheckedIOException;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class RconTest implements Closeable {
    private static final int SEND_BUFFER_SIZE = 8192;
    
    private final Socket socket;
    private final ExecutorService receiverExecutor;
    private CompletableFuture<String> previousFuture;
    private int nextId;
    
    private RconTest(String host, int port, String password, boolean nodelay) throws IOException {
        socket = new Socket(host, port);
        // Set buffer size to get consistent results across different devices
        socket.setSendBufferSize(SEND_BUFFER_SIZE);
        
        if (socket.getSendBufferSize() != SEND_BUFFER_SIZE) {
            throw new IOException("Setting buffer size failed");
        }
        
        receiverExecutor = Executors.newFixedThreadPool(1);
        previousFuture = CompletableFuture.completedFuture(null);
        
        sendMessage(3, 2, password, Runnable::run).join();
    }
    
    public void close() throws IOException {
        receiverExecutor.shutdown();
        socket.close();
    }
    
    private CompletableFuture<String> sendMessage(int type, int returnType, String message, Executor receiverExecutor) throws IOException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        DataOutputStream dataOut = new DataOutputStream(out);
        
        int id = nextId;
        byte[] bytes = message.getBytes(StandardCharsets.UTF_8);
        dataOut.writeInt(Integer.reverseBytes(bytes.length + 4 + 4 + 2));
        dataOut.writeInt(Integer.reverseBytes(id));
        dataOut.writeInt(Integer.reverseBytes(type));
        dataOut.write(bytes);
        dataOut.write(0);
        dataOut.write(0);
        
        dataOut.flush();
        socket.getOutputStream().write(out.toByteArray());
        
        nextId++;
        
        return previousFuture = previousFuture.thenApplyAsync(
            prev -> {
                try {
                    return receive(id, returnType);
                }
                catch (IOException ioException) {
                    throw new UncheckedIOException(ioException);
                }
            },
            receiverExecutor
        );
    }
    
    private String receive(int expectedId, int expectedType) throws IOException {
        InputStream in = socket.getInputStream();
        DataInputStream dataIn = new DataInputStream(in);
        
        int length = Integer.reverseBytes(dataIn.readInt());
        int id = Integer.reverseBytes(dataIn.readInt());
        int type = Integer.reverseBytes(dataIn.readInt());
        
        if (id != expectedId) {
            throw new IOException("Expected id " + expectedId + ", but got " + id);
        }
        if (type != expectedType) {
            throw new IOException("Expected type " + expectedType + ", but got " + type);
        }
        
        length -= 8 + 2; // 8 = id + type; 2 = trailing 0s
        
        byte[] responseBytes = new byte[length];
        dataIn.readFully(responseBytes);
        
        // Skip trailing two 0s
        int toSkip = 2;
        while (toSkip > 0) {
            toSkip -= dataIn.skipBytes(toSkip);
        }
        
        return new String(responseBytes, StandardCharsets.UTF_8);
    }
    
    public CompletableFuture<String> sendCommand(String command) throws IOException {
        return sendMessage(2, 0, command, receiverExecutor);
    }
    
    public static void main(String[] args) throws IOException {
        if (args.length != 3) {
            System.err.println("Arguments: <host> <port> <password>");
            System.exit(-1);
        }
        
        String host = args[0];
        int port = Integer.parseInt(args[1]);
        String password = args[2];
        
        int maxPacketLength = 1460;
        // length + id + type
        int packetMetaSize = 4 + 4 + 4;
        String commandPrefix = "execute run ".repeat(115) + "scoreboard players add ";
        String commandSuffix = " test 1";
        int commandBaseSize = commandPrefix.length() + commandSuffix.length() + 2; // + 2 for trailing 0s
        
        String command = commandPrefix + "a".repeat(maxPacketLength - packetMetaSize - commandBaseSize) + commandSuffix;
        
        try (RconTest rcon = new RconTest(host, port, password, false)) {
            System.out.println(rcon.sendCommand("scoreboard objectives add test dummy").join());
            
            for (int i = 0; i < 100; i++) {
                rcon.sendCommand(command).whenComplete((s, e) -> {
                    if (e != null) {
                        e.printStackTrace();
                    }
                    else {
                        System.out.println(s);
                    }
                });
            }
        }
    }
}