#define _GNU_SOURCE

#include <sys/timerfd.h>
#include <sys/epoll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <inttypes.h>
#include <linux/tcp.h>

enum event_id {
  TIMER,
  DATA_IN,
};

const int loop_period_ns = 100 * 1000;
#define MAX_EPOLL_EVENTS 2

int server(struct addrinfo *result) {
  int ret, msg_len, lfd, fd, one = 1, tfd, efd, n;
  struct addrinfo *rp;
  struct sockaddr addr = { 0 };
  socklen_t addrlen = sizeof(addr);
  struct timespec ts = { 0 };
  struct itimerspec timer = {
    .it_interval = { .tv_nsec = loop_period_ns },
    .it_value = { },
  };
  struct epoll_event events[MAX_EPOLL_EVENTS] = { { 0 } };
  int i, j;
  char msg[512];
  uint64_t timer_buf;
  struct tcp_info tcp_info = { 0 };
  socklen_t tcp_info_len = sizeof(tcp_info);

  for (rp = result; rp != NULL; rp = rp->ai_next) {
    lfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
    if (lfd != -1) {
      if (bind(lfd, rp->ai_addr, rp->ai_addrlen) != 0) {
        fprintf(stderr, "bind() failed: %m.\n");
        close(lfd);
      }
      else {
        break;
      }
    }
    else {
      fprintf(stderr, "socket() failed: %m.\n");
    }
  }
  
  if (rp == NULL) {
    fprintf(stderr, "Could not establish server.\n");
    return -1;
  }

  freeaddrinfo(result);

  ret = setsockopt(lfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
  if (ret == -1) {
    fprintf(stderr, "Could not setsockopt(SO_REUSEADDR): %m.\n");
    return -1;
  }

  ret = listen(lfd, 0);
  if (ret == -1) {
    fprintf(stderr, "Could not listen(): %m.\n");
    return -1;
  }

  fd = accept4(lfd, &addr, &addrlen, SOCK_NONBLOCK | SOCK_CLOEXEC);
  if (fd == -1) {
    fprintf(stderr, "Could not accept(): %m.\n");
    return -1;
  }

  tfd = timerfd_create(CLOCK_MONOTONIC, TFD_CLOEXEC);
  if (ret == -1) {
    fprintf(stderr, "Could not timerfd_create(): %m.\n");
    return -1;
  }

  ret = clock_gettime(CLOCK_MONOTONIC, &ts);
  if (ret == -1) {
    fprintf(stderr, "Could not clock_gettime(): %m.\n");
    return -1;
  }

  timer.it_value.tv_sec = ts.tv_sec + 1;
  ret = timerfd_settime(tfd, TFD_TIMER_ABSTIME, &timer, NULL);
  if (ret == -1) {
    fprintf(stderr, "Could not timerfd_settime(): %m.\n");
    return -1;
  }

  efd = epoll_create1(EPOLL_CLOEXEC);
  if (efd == -1) {
    fprintf(stderr, "Could not epoll_create1(): %m.\n");
    return -1;
  }

  events[0].events = EPOLLIN | EPOLLRDHUP;
  events[0].data.u32 = TIMER;
  ret = epoll_ctl(efd, EPOLL_CTL_ADD, tfd, &(events[0]));
  if (ret == -1) {
    fprintf(stderr, "Could not epoll_ctl(EPOLL_CTL_ADD, tfd): %m.\n");
    return -1;
  }

  events[0].events = EPOLLIN | EPOLLRDHUP;
  events[0].data.u32 = DATA_IN;
  ret = epoll_ctl(efd, EPOLL_CTL_ADD, fd, &(events[0]));
  if (ret == -1) {
    fprintf(stderr, "Could not epoll_ctl(EPOLL_CTL_ADD, tfd): %m.\n");
    return -1;
  }

  uint64_t bytes_acked = 0;
  int data_received = 0;
  int received_data, received_ack;

  j = 0;
  /* TODO: handle ERR, HUP and RDHUP for all file descriptor kinds. */
  while ((n = epoll_wait(efd, events, MAX_EPOLL_EVENTS, -1))) {
    ret = clock_gettime(CLOCK_MONOTONIC, &ts);
    if (ret == -1) {
      fprintf(stderr, "Could not clock_gettime(): %m.\n");
      return -1;
    }

    received_data = 0;
    received_ack = 0;

    for (i = 0; i < n; i++) {
      switch (events[i].data.u32) {
        case TIMER:
          j++;
          ret = read(tfd, &timer_buf, sizeof(timer_buf));
          if (ret == -1) {
            fprintf(stderr, "Could not read timer infos: %m.\n");
            return -1;
          }

          ret = getsockopt(fd, IPPROTO_TCP, TCP_INFO, &tcp_info, &tcp_info_len);
          if (ret == -1) {
            fprintf(stderr, "Could not get struct tcp_info: %m.\n");
            return -1;
          }

          if (tcp_info.tcpi_bytes_acked != bytes_acked) {
            bytes_acked = tcp_info.tcpi_bytes_acked;
            received_ack = 1;
          }

          if ((j % (1000 * 1000 * 1000 / loop_period_ns))  == 0) {
            j = 0;
            ret = snprintf(msg, sizeof(msg),
                "%" PRIu64 ".%09ld: %03d\n",
                ts.tv_sec, ts.tv_nsec, j);
            if ((ret < 0) || ((size_t) ret >= sizeof(msg))) {
              fprintf(stderr,
                  "Could not write message in full: needed %d bytes.",
                  ret);
              return -1;
            }
            msg_len = ret;

            fprintf(stdout, "Sending message: %s", msg);
            fflush(stdout);
            errno = 0;
            ret = send(fd, msg, msg_len, 0);
            if (ret != msg_len) {
              fprintf(stderr, "Could not send message in full: %m.\n");
            }
          }
          break;
        case DATA_IN:
          ret = read(fd, &msg, sizeof(msg));
          if (ret == -1) {
            fprintf(stderr, "Could not read data: %m.\n");
            return -1;
          }

          data_received += ret;
          received_data = 1;

          j++;
          ret = read(tfd, &timer_buf, sizeof(timer_buf));
          if (ret == -1) {
            fprintf(stderr, "Could not read timer infos: %m.\n");
            return -1;
          }

          ret = getsockopt(fd, IPPROTO_TCP, TCP_INFO, &tcp_info, &tcp_info_len);
          if (ret == -1) {
            fprintf(stderr, "Could not get struct tcp_info: %m.\n");
            return -1;
          }

          if (tcp_info.tcpi_bytes_acked != bytes_acked) {
            bytes_acked = tcp_info.tcpi_bytes_acked;
            received_ack = 1;
          }

          goto next_events;

          break;
        default:
          fprintf(stderr, "Unknown event kind %d, skipping.\n",
              events[i].data.u32);
          break;
      }
    }

next_events:

    if (received_ack || received_data) {
      fprintf(stdout,
          "%" PRIu64 ".%09ld %06" PRIu64 " %06d %s%s\n",
          ts.tv_sec, ts.tv_nsec,
          bytes_acked,
          data_received,
          (received_ack ? "X" : "_"),
          (received_data ? "X" : "_")
          );
      fflush(stdout);
    }

  }

  return 0;
}

int client(struct addrinfo *result) {
  int ret, ret2, fd;
  struct addrinfo *rp;
  char msg[512];
  
  for (rp = result; rp != NULL; rp = rp->ai_next) {
    fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
    if (fd != -1) {
      if (connect(fd, rp->ai_addr, rp->ai_addrlen) != 0) {
        fprintf(stderr, "connect() failed: %m.\n");
        close(fd);
      }
      else {
        break;
      }
    }
    else {
      fprintf(stderr, "socket() failed: %m.\n");
    }
  }
  
  if (rp == NULL) {
    fprintf(stderr, "Could not establish client.\n");
    return -1;
  }

  freeaddrinfo(result);

  while (1) {
    ret = read(fd, msg, sizeof(msg));
    if (ret == -1) {
      fprintf(stderr, "Could not read data: %m.\n");
      return -1;
    }

    errno = 0;
    ret2 = write(fd, msg, ret);
    if (ret2 != ret) {
      fprintf(stderr, "Could not echo data: %m [%d, %d].\n", ret, ret2);
      return -1;
    }
  }

  return 0;
}

int main(int argc, char *argv[]) {
  int ret;
  struct addrinfo *result;
  struct addrinfo settings = {
    .ai_family = AF_UNSPEC,
    .ai_socktype = SOCK_STREAM,
    .ai_flags = AI_PASSIVE,
  };

  if ((argc != 4)
      || (strcmp(argv[1], "server") != 0 && (strcmp(argv[1], "client") != 0)))
  {
    fprintf(stderr, "Usage: %s <client|server> <IP> <port>.\n", argv[0]);
    return -1;
  }

  ret = getaddrinfo(argv[2], argv[3], &settings, &result);
  if (ret != 0) {
    fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(ret));
    return -1;
  }

  switch(argv[1][0]) {
    case 's':
      return server(result);
    case 'c':
      return client(result);
  }
}