#include <arpa/inet.h>
#include <errno.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

#define PORT 1234
#define NUM_CONNECTIONS 12
#define DATA_SIZE (4 * 1024)

static int read_full(int fd, void *buf, size_t len) {
  static const size_t sizes[] = {1, 7, 17, 255, 256, 257, 511, 512, 4096};
  char *p = (char *) buf;
  size_t i = 0;
  while (len > 0) {
    size_t nreq = sizes[i++ % (sizeof(sizes) / sizeof(sizes[0]))];
    ssize_t n;
    if (nreq > len) nreq = len;
    n = recv(fd, p, nreq, 0);
    if (n <= 0) return -1;
    p += n;
    len -= (size_t) n;
  }
  return 0;
}

static int write_full(int fd, const void *buf, size_t len) {
  const char *p = (const char *) buf;
  while (len > 0) {
    ssize_t n = send(fd, p, len, 0);
    if (n <= 0) return -1;
    p += n;
    len -= (size_t) n;
  }
  return 0;
}

static void fill_text(char *buf, size_t len) {
  size_t i;
  for (i = 0; i < len; i++) buf[i] = (char) (' ' + random() % 95);
}

static void *worker(void *arg) {
  const char *host = (const char *) arg;
  char tx[DATA_SIZE], rx[DATA_SIZE];
  struct sockaddr_in sa;
  struct timeval timeout;
  unsigned port = 0;
  int fd;

  fill_text(tx, sizeof(tx));
  memset(rx, 0, sizeof(rx));

  fd = socket(AF_INET, SOCK_STREAM, 0);
  if (fd < 0) goto fail;
  timeout.tv_sec = 5;
  timeout.tv_usec = 0;
  if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) != 0 ||
      setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) != 0)
    goto fail_close;

  memset(&sa, 0, sizeof(sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons(PORT);
  if (inet_pton(AF_INET, host, &sa.sin_addr) != 1) goto fail_close;
  printf("%d CONNECTING\n", fd);
  if (connect(fd, (struct sockaddr *) &sa, sizeof(sa)) != 0) goto fail_close;
  {
    socklen_t len = sizeof(sa);
    if (getsockname(fd, (struct sockaddr *) &sa, &len) != 0) goto fail_close;
    port = (unsigned) ntohs(sa.sin_port);
  }
  printf("%d CONNECTED port: %u\n", fd, port);
  printf("%d WRITING\n", fd);
  if (write_full(fd, tx, sizeof(tx)) != 0) goto fail_close;
  printf("%d READING\n", fd);
  if (read_full(fd, rx, sizeof(rx)) != 0) goto fail_close;
  if (memcmp(tx, rx, sizeof(tx)) != 0) goto fail_close;

  printf("%d OK-CLOSING\n", fd);
  close(fd);
  return NULL;

fail_close:
  printf("%d ERROR-CLOSING\n", fd);
  close(fd);
fail:
  exit(1);
}

enum wait_mode { WAIT_ABORT, WAIT_DRAIN, WAIT_HOLD, WAIT_CLOSE };

// Wait for output, optionally draining through the final marker.
static int wait_data(const char *host, enum wait_mode mode) {
  struct sockaddr_in sa;
  struct timeval timeout;
  char buf[512];
  ssize_t n;
  size_t i;
  int fd = socket(AF_INET, SOCK_STREAM, 0);

  if (fd < 0) return 1;
  memset(&sa, 0, sizeof(sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons(PORT);
  if (inet_pton(AF_INET, host, &sa.sin_addr) != 1) goto fail;
  timeout.tv_sec = 5, timeout.tv_usec = 0;
  if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) != 0)
    goto fail;
  if (connect(fd, (struct sockaddr *) &sa, sizeof(sa)) != 0) goto fail;
  if (mode == WAIT_HOLD) {
    usleep(100000);  // Keep the peer connected while the server fails TX
    close(fd);
    return 0;
  }
  if (mode == WAIT_CLOSE) {
    n = recv(fd, buf, sizeof(buf), 0);
    close(fd);
    return n == 0 || (n < 0 && errno == ECONNRESET) ? 0 : 1;
  }
  for (;;) {
    n = recv(fd, buf, sizeof(buf), 0);
    if (n <= 0) goto fail;
    if (mode == WAIT_ABORT) break;
    for (i = 0; i < (size_t) n; i++) {
      if (buf[i] == '\1') break;
    }
    if (i < (size_t) n) break;
  }
  close(fd);
  return 0;

fail:
  close(fd);
  return 1;
}

int main(int argc, char *argv[]) {
  pthread_t threads[NUM_CONNECTIONS];
  int connections = NUM_CONNECTIONS, i;

  if (argc == 3 && strcmp(argv[2], "waitandabort") == 0)
    return wait_data(argv[1], WAIT_ABORT);
  if (argc == 3 && strcmp(argv[2], "waitanddrain") == 0)
    return wait_data(argv[1], WAIT_DRAIN);
  if (argc == 3 && strcmp(argv[2], "hold") == 0)
    return wait_data(argv[1], WAIT_HOLD);
  if (argc == 3 && strcmp(argv[2], "expectclose") == 0)
    return wait_data(argv[1], WAIT_CLOSE);
  if (argc == 3) connections = atoi(argv[2]);
  if ((argc != 2 && argc != 3) || connections < 1 ||
      connections > NUM_CONNECTIONS) {
    fprintf(stderr, "usage: %s IP [connections|waitandabort|waitanddrain|hold|expectclose]\n",
            argv[0]);
    return 1;
  }

  srandom(1);

  for (i = 0; i < connections; i++) {
    if (pthread_create(&threads[i], NULL, worker, argv[1]) != 0) return 1;
  }

  for (i = 0; i < connections; i++) {
    if (pthread_join(threads[i], NULL) != 0) return 1;
  }

  return 0;
}