package dev.elaine.network;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;

public final class NioEchoServer implements AutoCloseable {
    private final Selector selector;
    private final ServerSocketChannel server;
    private final Thread loopThread;
    private volatile boolean running = true;

    public NioEchoServer() throws IOException {
        selector = Selector.open();
        server = ServerSocketChannel.open();
        server.configureBlocking(false);
        server.bind(new InetSocketAddress("127.0.0.1", 0));
        server.register(selector, SelectionKey.OP_ACCEPT);
        loopThread = Thread.ofPlatform().name("nio-event-loop").start(this::runLoop);
    }

    public int port() throws IOException {
        return ((InetSocketAddress) server.getLocalAddress()).getPort();
    }

    private void runLoop() {
        try {
            while (running) {
                selector.select();
                Iterator<SelectionKey> keys = selector.selectedKeys().iterator();
                while (keys.hasNext()) {
                    SelectionKey key = keys.next();
                    keys.remove();
                    if (key.isValid() && key.isAcceptable()) {
                        accept();
                    } else if (key.isValid() && key.isReadable()) {
                        echo(key);
                    }
                }
            }
        } catch (IOException error) {
            if (running) {
                throw new IllegalStateException("NIO 事件循环失败", error);
            }
        }
    }

    private void accept() throws IOException {
        SocketChannel channel = server.accept();
        if (channel != null) {
            channel.configureBlocking(false);
            channel.register(selector, SelectionKey.OP_READ, ByteBuffer.allocateDirect(1024));
        }
    }

    private static void echo(SelectionKey key) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        ByteBuffer buffer = (ByteBuffer) key.attachment();
        int count = channel.read(buffer);
        if (count < 0) {
            channel.close();
            return;
        }
        buffer.flip();
        while (buffer.hasRemaining()) {
            channel.write(buffer);
        }
        buffer.clear();
    }

    @Override
    public void close() throws Exception {
        running = false;
        selector.wakeup();
        loopThread.join();
        for (SelectionKey key : selector.keys()) {
            key.channel().close();
        }
        selector.close();
    }

    public static void main(String[] args) throws Exception {
        try (NioEchoServer server = new NioEchoServer();
             SocketChannel client = SocketChannel.open(
                     new InetSocketAddress("127.0.0.1", server.port()))) {
            client.write(StandardCharsets.UTF_8.encode("你好，NIO"));
            ByteBuffer reply = ByteBuffer.allocate(32);
            client.read(reply);
            reply.flip();
            System.out.println(StandardCharsets.UTF_8.decode(reply));
        }
    }
}