#include "test.h"
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#define TEST(c)                                                                \
  do {                                                                         \
    errno = 0;                                                                 \
    if (!(c))                                                                  \
      t_error("%s failed (errno = %d)\n", #c, errno);                          \
  } while (0)

#define N 10

int main() {
  int listener_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);

  // Setup a listener bound to port 0 to have the OS assign us one.
  struct sockaddr_in server_address;
  socklen_t server_address_len = sizeof(server_address);
  server_address.sin_family = AF_INET;
  server_address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  server_address.sin_port = 0;
  TEST(bind(listener_fd, (struct sockaddr *)&server_address,
            sizeof(server_address)) != -1);
  TEST(getsockname(listener_fd, (struct sockaddr *)&server_address,
                   &server_address_len) != -1);
  TEST(listen(listener_fd, N) != -1);

  struct pollfd clients[N];
  int remaining = N;

  // Connect N clients, then wait for them all to reach the writable state
  // meaning that they're more-or-less
  // connected.
  for (int i = 0; i < N; i++) {
    int client_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
    TEST(client_fd != -1);
    int rc;
    TEST((rc = connect(client_fd, (struct sockaddr *)&server_address,
                       server_address_len)) != -1 ||
         errno == EINPROGRESS);
    clients[i].fd = client_fd;
    if (rc == -1) {
      clients[i].events = POLLWRNORM;
    } else {
      remaining--;
      clients[i].events = 0;
    }
  }

  while (remaining > 0) {
    TEST(poll(clients, N, -1) > 0);

    for (int i = 0; i < N; i++) {
      if (clients[i].revents & POLLWRNORM) {
        clients[i].events = 0;
        remaining--;
      }
    }
  }

  // Now accept everything quickly and close it all.
  for (int i = 0; i < N; i++) {
    struct pollfd listener;
    listener.fd = listener_fd;
    listener.events = POLLRDNORM;
    TEST(poll(&listener, 1, -1) == 1);
    int server_fd;
    TEST((server_fd = accept(listener_fd, NULL, NULL)) != -1);
    TEST(close(server_fd) != -1);
  }

  for (int i = 0; i < N; i++)
    TEST(close(clients[i].fd) != -1);

  return t_status;
}
