package two;

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.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class NIOServer {
    private Selector selector;
    private Map<SocketChannel, String> clients = new HashMap<>();

    public void start(int port) throws IOException {
        // 创建ServerSocketChannel
        ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.socket().bind(new InetSocketAddress(port));
        serverSocketChannel.configureBlocking(false);

        // 创建Selector
        selector = Selector.open();

        // 注册ServerSocketChannel到Selector上，监听ACCEPT事件
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

        System.out.println("Server started...");

        // 循环处理事件
        while (true) {
            // 调用select方法阻塞，等待事件发生
            selector.select();

            Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();
            while (keyIterator.hasNext()) {
                SelectionKey key = keyIterator.next();
                keyIterator.remove();

                if (key.isAcceptable()) { // 处理ACCEPT事件
                    handleAccept(key);
                } else if (key.isReadable()) { // 处理READ事件
                    handleRead(key);
                }
            }
        }
    }

    private void handleAccept(SelectionKey key) throws IOException {
        // 获取ServerSocketChannel
        ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key.channel();

        // 接受客户端连接
        SocketChannel socketChannel = serverSocketChannel.accept();
        socketChannel.configureBlocking(false);

        // 注册SocketChannel到Selector上，监听READ事件
        socketChannel.register(selector, SelectionKey.OP_READ);

        System.out.println("Client " + socketChannel.getRemoteAddress() + " connected.");

        // 将新连接的客户端添加到clients列表中
        clients.put(socketChannel, socketChannel.getRemoteAddress().toString());
    }

    private void handleRead(SelectionKey key) throws IOException {
        // 获取SocketChannel
        SocketChannel socketChannel = (SocketChannel) key.channel();

        // 读取数据
        ByteBuffer buffer = ByteBuffer.allocate(1024);
        int numRead = socketChannel.read(buffer);
        if (numRead == -1) { // 客户端关闭连接
            key.cancel();
            clients.remove(socketChannel);
            System.out.println("Client " + socketChannel.getRemoteAddress() + " disconnected.");
            socketChannel.close();
            return;
        }

        // 将数据转发给其他客户端
        String message = new String(buffer.array(), 0, numRead).trim();
        System.out.println("Received message from " + socketChannel.getRemoteAddress() + ": " + message);
        for (SocketChannel client : clients.keySet()) {
            if (client != socketChannel) {
                ByteBuffer writeBuffer = ByteBuffer.wrap(message.getBytes());
                client.write(writeBuffer);
            }
        }
    }

    public static void main(String[] args) throws IOException {
        NIOServer server = new NIOServer();
        server.start(8888);
    }
}