package info.gamlor.redis;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;

import static info.gamlor.redis.Writer.ensureAvailableWriteSpace;
import static java.nio.charset.StandardCharsets.UTF_8;

public class RedisMain {

    public static void main(String[] args) throws Exception {
        var scheduler = Executors.newVirtualThreadPerTaskExecutor();
        var socket = ServerSocketChannel.open();
        socket.bind(new InetSocketAddress("0.0.0.0", 16379));

        System.out.println("App is listening on 0.0.0.0:16379");
        var clone = new RedisClone();
        while (true) {
            var client = socket.accept();
            scheduler.execute(() -> {
                try (client) {
                    clone.handleConnection(client);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        }
    }
}

class RedisClone {
    private final ConcurrentHashMap<String, String> state = new ConcurrentHashMap<>();

    public void handleConnection(SocketChannel socket) throws Exception {
        var args = new ArrayList<String>();
        // Configure the channel to be non-blocking: Our Writer and Reader class will control the blocking
        socket.configureBlocking(false);
        // Replace the JDK IO streams with our Reader and Writer
        var writer = new Writer(socket);
        var reader = new Reader(socket, writer);

        while (true) {
            args.clear();
            var line = reader.readLine();
            if (line == null)
                break;

            if (line.charAt(0) != '*')
                throw new RuntimeException("Cannot understand arg batch: " + line);

            var argsv = Integer.parseInt(line.substring(1));
            for (int i = 0; i < argsv; i++) {
                line = reader.readLine();
                if (line == null || line.charAt(0) != '$')
                    throw new RuntimeException("Cannot understand arg length: " + line);
                var argLen = Integer.parseInt(line.substring(1));
                line = reader.readLine();
                if (line == null || line.length() != argLen)
                    throw new RuntimeException("Wrong arg length expected " + argLen + " got: " + line);

                args.add(line);
            }

            var reply = executeCommand(args);
            if (reply == null) {
                writer.write("$-1\r\n");
            } else {
                writer.write("$" + reply.length() + "\r\n" + reply + "\r\n");
            }
        }
    }

    String executeCommand(List<String> args) {
        switch (args.get(0)) {
            case "GET":
                return state.get(args.get(1));
            case "SET":
                state.put(args.get(1), args.get(2));
                return null;
            default:
                throw new IllegalArgumentException("Unknown command: " + args.get(1));
        }
    }
}

class Writer {
    final SocketChannel socket;
    final ByteBuf writeBuffer = Unpooled.buffer(4 * 1024);

    public Writer(SocketChannel socket) throws IOException {
        this.socket = socket;
        assert !socket.isBlocking();
    }

    public void write(String toWrite) throws Exception {
        var bytes = toWrite.getBytes(UTF_8);
        ensureAvailableWriteSpace(writeBuffer);
        writeBuffer.writeBytes(bytes);

        final var AUTO_FLUSH_LIMIT = 1024;
        if (AUTO_FLUSH_LIMIT < writeBuffer.readableBytes()) {
            // A bit confusing in this use case: We read the buffers content into the socket: aka write to the socket
            var written = writeBuffer.readBytes(socket, writeBuffer.readableBytes());
            // If we want proper handling of the back pressure by waiting for the channel to be writable.
            // But for this example we ignore such concerns and just grow the writeBuffer defiantly
        }
    }

    public void flush() throws IOException {
        if (writeBuffer.readableBytes() > 0) {
            writeBuffer.readBytes(socket, writeBuffer.readableBytes());
        }
    }

    // The Netty ByteBufs are not circular buffer: Writes always go to the end and may grow the buffer
    // I assume the underlying reason is to make it more efficient to interact with Java NIO.
    // So, if we're running out of writeable space, discard the bytes already written and
    // copy the not yet read bytes to the start of the buffer, giving it enough space to write more at the end.
    static int ensureAvailableWriteSpace(ByteBuf buf) {
        final var MIN_WRITE_SPACE = 1024;

        if (buf.writableBytes() < MIN_WRITE_SPACE) {
            buf.discardReadBytes();
        }
        return Math.max(MIN_WRITE_SPACE, buf.writableBytes());
    }
}

class Reader {
    final SocketChannel socket;
    final Writer writer;
    final ByteBuf readBuffer = Unpooled.buffer(8 * 1024);
    private final Selector selector;


    public Reader(SocketChannel socket, Writer writer) throws IOException {
        this.socket = socket;
        this.writer = writer;
        this.selector = Selector.open();
        socket.register(selector, SelectionKey.OP_READ, this);
    }

    public String readLine() throws Exception {
        var eof = false;
        while (!eof) {
            var readIndex = readBuffer.readerIndex();
            var toIndex = readBuffer.readableBytes();
            // Find the next line in the read content
            var foundNewLine = readBuffer.indexOf(readIndex, readIndex + toIndex, (byte) '\n');
            if (foundNewLine >= 0) {
                var length = foundNewLine - readIndex;
                String line = readBuffer.toString(readIndex, length - 1, UTF_8);
                readBuffer.readerIndex(readIndex + length + 1);
                return line;
            } else {
                // Otherwise, read from the socket
                int readSize = ensureAvailableWriteSpace(readBuffer);
                // A bit confusing in this use case: We write the content of the socket into the buffer: aka read from the channel
                var read = readBuffer.writeBytes(socket, readSize);
                if (read < 0) {
                    eof = true;
                } else if (read == 0) {
                    // If we read nothing, ensure we flushed our previous reponses
                    writer.flush();
                    // And then wait until the socket becomes readable again
                    selector.select(key -> {
                        if (!key.isReadable()) {
                            throw new AssertionError("Expect to be readable again");
                        }
                    });
                }
            }
        }
        return null;
    }

}