!7

import java.nio.*;
import java.nio.channels.*;

static long timeout = 30000+10000; // idle time + latency
static int port = 6000;
static int maxLineLength = 2048;
static ConnectedInstances cI;
static ReliableSingleThread rst = new(r update);
static int theNumber;
static Map<SocketChannel, State> dataMapper = synchroMap();
static ServerSocketChannel serverChannel;
static Selector selector;
static Map<S, Channel> channels = synchroMap();
static Lock channelsLock = lock();

sclass State {
  S socketID = aGlobalID(), computerID;
  int countSeen = -1;
  new LineBuffer lineBuffer;
  long lastMessage = sysNow();
  Map<S, Channel> subbedChannels = synchroLinkedHashMap();
}

sclass Channel<A> {
  Set<SocketChannel> sockets = synchroSet();
  Map<SocketChannel, A> dataMapper = synchroMap();
  Lock lock = lock();
  
  void addSocket(SocketChannel socket) {
    lock lock;
    sockets.add(socket);
    initialMessage(socket);
  }
  
  void removeSocket(SocketChannel socket) {
    lock lock;
    sockets.remove(socket);
  }
  
  void update() {}
  void initialMessage(SocketChannel socket) {}
  void handleLine(S line) {}
}

sclass ComputerIDsChannel extends Channel<O> {
  L<S> value;
  
  void initialMessage(SocketChannel socket) {
    lock lock;
    sendTo(socket);
  }
  
  void update() {
    if (set_trueIfChanged(this, value := asList(cI.set())))
      for (SocketChannel s : sockets) sendTo(s);
  }
  
  void sendTo(SocketChannel socket) {
    sendLineToSocket(socket, "ComputerIDs = " + struct(value));
  }
}

sclass PublicChatChannel extends Channel<O> {
  L<S> msgs = synchroLinkedList();
  
  void handleLine(S line) {
    postMessage(unquote(line));
  }
  
  void postMessage(S msg) {
    msgs.add(msg);
  }
  
  void update {
    S msg;
    while ((msg = syncPopFirst(msgs)) != null) {
      S quoted = quote(msg);
      print("Chat send: " + shorten(quoted, 20));
      quoted = "chat:" + quoted;
      for (SocketChannel s : sockets)
        sendLineToSocket(s, quoted);
    }
  }
}

p {
  cI = new ConnectedInstances;
  cI.setAfterglowInterval(1000);
  cI.connected2.onChange(rst);
  cI.start();
  
  selector = Selector.open();
	serverChannel = ServerSocketChannel.open();
	serverChannel.configureBlocking(false);
	
  InetSocketAddress sockAddr = new InetSocketAddress(port);
	serverChannel.socket().bind(sockAddr);
	serverChannel.register(selector, SelectionKey.OP_ACCEPT);
	
	thread "Server Accept Loop" { acceptLoop(); }
	
	doEvery(1000, r {
	  for (SocketChannel socket : keysList(dataMapper)) {
	    State state = dataMapper.get(socket);
	    if (state != null && sysTime() >= state.lastMessage+timeout) pcall {
	      socket.close();
	      onSocketClose(socket);
	    }
	  }
	  for (Channel channel : values(channels)) pcall {
	    lock channel.lock;
      channel.update();
    }
	});
}

svoid cleanMeUp() ctex {
  if (serverChannel != null)
    serverChannel.close();
  for (SocketChannel socket : keys(getAndClearMap(dataMapper))) pcall {
    socket.close();
  }
}

svoid acceptLoop ctex {
  while licensed {
    // wait for events
    selector.select();

	  Iterator<SelectionKey> keys = selector.selectedKeys().iterator();
	  while (keys.hasNext()) {
      SelectionKey key = keys.next();
	    keys.remove();
	    
	    if (!key.isValid()) continue;
	    pcall {
        if (key.isAcceptable())
          acceptOnSocket(key);
	      else if (key.isReadable())
          readFromSocket(key);
      }
	  }
	}
}

//accept a connection made to this channel's socket
static void acceptOnSocket(SelectionKey key) throws IOException {
  ServerSocketChannel serverChannel = cast key.channel();
  SocketChannel channel = serverChannel.accept();
  channel.configureBlocking(false);
  Socket actualSocket = channel.socket();
  SocketAddress remoteAddr = actualSocket.getRemoteSocketAddress();
  print("Connected to: " + remoteAddr);
  
  // register channel with selector for further IO
	dataMapper.put(channel, new State);
	channel.register(selector, SelectionKey.OP_READ);
}

svoid onSocketClose(SocketChannel socket) {
  Socket actualSocket = socket.socket();
  SocketAddress remoteAddr = actualSocket.getRemoteSocketAddress();
  print("Connection closed by client: " + remoteAddr);
  State state = dataMapper.get(socket);
  if (state != null && state.computerID != null)
    cI.lostConnection(state.computerID);
  dataMapper.remove(socket);
  removeFromAllChannels(socket);
}

static void readFromSocket(SelectionKey key) throws IOException {
	final SocketChannel socket = cast key.channel();
  ByteBuffer buffer = ByteBuffer.allocate(1024);
  int numRead = socket.read(buffer);
  if (numRead == -1) {
    socket.close();
    onSocketClose(socket);
    key.cancel();
    ret;
  }

  S s = new String(buffer.array(), 0, numRead);
  //print("Got data from client: " + s);
  final State state = dataMapper.get(socket);
  if (state != null) {
    state.lastMessage = sysNow();
    state.lineBuffer.append(s, voidfunc(S line) {
      if (nempty(line)) print("Got line: " + line);
      pcall { handleLine(line, socket, state); }
    });
    if (state.lineBuffer.size() > maxLineLength) {
      print("Bad client (line too long");
      socket.close();
    }
  }
}

svoid handleLine(S line, SocketChannel socket, State state) {
  new Matches m;
  if (empty(line)) ret with sendLineToSocket(socket, ""); // idle
  
  if (jmatch("computerID=*", line, m)) {
    if (state.computerID != null || !possibleComputerID($1)) ret;
    print("Got computer ID: " + $1);
    state.computerID = $1;
    cI.gotConnection(state.computerID);
    ret;
  }
  
  if (jmatch("sub *", line, m)) {
    Channel channel = getChannelForSocket(socket, $1);
    if (channel == null)
      sendLineToSocket(socket, "Channel " + $1 + " not found");
    else
      sendLineToSocket(socket, "Subbed " + $1);
    ret;
  }
  
  L<S> tok = javaTok(line);
  if (eq(get(tok, 3), ":")) {
    Channel channel = state.subbedChannels.get(get(tok, 1));
    if (channel != null)
      channel.handleLine(joinSubList(tok, 5));
  }
  
  // ignore unknown line
}

html { ret "Use port " + port; }

svoid sendNumber(SocketChannel socket) ctex {
  State s = dataMapper.get(socket);
  if (s != null && theNumber != s.countSeen) {
    s.countSeen = theNumber;
    print("Sending number to " + or2(s.computerID, "?") + ": " + theNumber);
    sendLineToSocket(socket, str(theNumber));
  }
}

svoid sendLineToSocket(SocketChannel socket, S s) ctex {
  socket.write(ByteBuffer.wrap(toUtf8(s + "\n")));
}

svoid removeFromAllChannels(SocketChannel socket) {
  lock channelsLock;
  for (S name, Channel c : cloneMap(channels)) { pcall {
    c.removeSocket(socket);
    if (empty(c.sockets)) {
      print("Deleting channel: " + name);
      channels.remove(name);
    }
  }}
}

static Channel getChannelForSocket(SocketChannel socket, S name) {
  lock channelsLock;
  Channel c = channels.get(name);
  if (c == null) {
    c = makeChannel(name);
    if (c == null) null;
    print("New channel: " + name);
    channels.put(name, c);
    lock c.lock;
    c.update();
  }
  c.addSocket(socket);
  dataMapper.get(socket).subbedChannels.put(name, c);
  ret c;
}

static Channel makeChannel(S name) {
  if (eq(name, "computerIDs")) ret new ComputerIDsChannel;
  if (eq(name, "chat")) ret new PublicChatChannel;
    
  ret null with print("Unknown channel: " + name);
}

svoid update {
  print("OS Instances update");
  if (empty(dataMapper)) ret;

  int value = cI.value();
  //if (value == theNumber) ret;
  theNumber = value;
  for (SocketChannel socket : keysList(dataMapper)) pcall {
    sendNumber(socket);
  }
  print("OS Instances update done");
}