// Small sockets example to investigate issue on Barrelfish, where client
// blocks if both ends run in the same domain.

#include <errno.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#ifdef __barrelfish__
# include <barrelfish/threads.h>
#else
# include <pthread.h>
#endif

#ifdef __barrelfish__
#include <dmalloc/dmalloc.h>     // for dlmalloc, dlrealloc, dlfree.
#include <barrelfish/debug.h>    // for USER_PANIC.
#include <barrelfish/threads.h>  // for struct thread, thread_create.
#include <barrelfish/waitset.h>  // for struct waitset, waitset_init.
#include <lwip/sockets.h>  // for lwip_socket_init.
#include <lwip/tcpip.h>    // for tcpip_init.
#include <vfs/vfs.h>       // for vfs_init.

// Initialize networking subsystem.
// (Taken from {BARRELFISH}/usr/net-test/net-test.c)

// Implemented in {BARRELFISH}/lib/lwip/src/barrelfish/idc_net_control.c.
extern void network_polling_loop(void);

static struct thread_mutex network_setup_mutex = THREAD_MUTEX_INITIALIZER;
static struct waitset lwip_waitset;
static struct thread *network_poll_thread;

static int poll_loop(void *args) {
  (void) args;  // unused
  network_polling_loop();
  return 0;
}

static void network_setup_helper(void) {
  waitset_init(&lwip_waitset);
  thread_mutex_lock(&network_setup_mutex);
  // FIXME: replaced lwip_init_auto_ex with lwip_init_auto without testing
  //    lwip_init_auto_ex(&lwip_waitset, &network_setup_mutex);

  // lwip_init_auto();
  tcpip_init(NULL, NULL); // FIXME: serious problem: why is it called after
  // calling lwip_init_auto? henc lwip_init_auto is
  // commented out (without testing!!!)
  lwip_socket_init();
  thread_mutex_unlock(&network_setup_mutex);

  // start an event dispatch thread
  network_poll_thread = thread_create(poll_loop, NULL);
  if (network_poll_thread == NULL) {
    USER_PANIC("thread_create() returned NULL.");
  }
}

// Enable Doug Lea's malloc.
typedef void *(*alt_malloc_t)(size_t bytes);
extern alt_malloc_t alt_malloc;
typedef void (*alt_free_t)(void *p);
extern alt_free_t alt_free;
typedef void *(*alt_realloc_t)(void *p, size_t bytes);
extern alt_realloc_t alt_realloc;

static void dmalloc_init(void) {
  alt_malloc = &dlmalloc;
  alt_free = &dlfree;
  alt_realloc = &dlrealloc;
}

// This bit of trickery is required to ensure the VFS libc glue is initialized
// before global/static constructors have run.
//
// Notes:
// * This method is a little fragile, but seems to work in practice with a
//   GCC 4.7.2 cross-compiler.
// * By using this GCC attribute, this function will be added to the .ctors
//   section (hopefully before any other C++ constructors!).
// * If this is approach is problematic, we could instead use
//     __attribute__((section(".preinit_array")))
//   but that will require some changes to {BARRELFISH}/lib/crt/crtbegin.c.
// * For more details see: http://www.acsu.buffalo.edu/~charngda/elf.html
static void initializeBeforeMain() __attribute__((constructor(101)));
static void initializeBeforeMain() {
  dmalloc_init();
  vfs_init();
  network_setup_helper();
}
#endif // __barrelfish__

static int kServerPort = 5000;

#ifdef __barrelfish__
// Need to hardcode IP address here since there is no concept of 'localhost'.
static const char* kServerIpAddress = "10.110.4.21";      // gruyere
#else
static const char* kServerIpAddress = "127.0.0.1";
#endif

static volatile bool server_is_ready = false;

static void perror_and_exit(const char* msg) {
  perror(msg);
  exit(1);
}

static void run_server(void) {
  struct sockaddr_in address = { 0 };
  address.sin_family = AF_INET;
  address.sin_addr.s_addr = htonl(INADDR_ANY);
  address.sin_port = htons(kServerPort);

  int listenfd = socket(AF_INET, SOCK_STREAM, 0);
  if (listenfd < 0)
    perror_and_exit("server socket");

#ifndef __barrelfish__
  const int on = true;
  setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
#endif
  int retcode = bind(listenfd, (struct sockaddr*)&address, sizeof(address));
  if (retcode < 0)
    perror_and_exit("bind");

  retcode = listen(listenfd, 10);
  if (retcode < 0)
    perror_and_exit("listen");

  printf("[server] listening on port %d.\n", kServerPort);
  server_is_ready = true;

  while (true) {
    int connfd = accept(listenfd, /*addr*/ NULL, /*addrlen*/ NULL);
    if (connfd < 0) {
      perror_and_exit("accept");
    }
    printf("[server] accepted connection.\n");

    char send_buffer[1025] = { '\0' };
    time_t ticks = time(NULL);
    snprintf(send_buffer, sizeof(send_buffer), "The time is %s", ctime(&ticks));
    printf("[server] calling write on socket (buffer=%p, length=%zu)\n",
           send_buffer, strlen(send_buffer));
    int bytes_written = write(connfd, send_buffer, strlen(send_buffer));
    if (bytes_written < 0)
      perror_and_exit("write");
    printf("[server] sent reply (%d bytes).\n", bytes_written);

    close(connfd);
    //sleep(1);
    //printf("[server] back from snooze ...\n");

    break; // just once for testing purposes.
  }

  printf("[server] DONE\n");
}

static void run_client(void) {
  while (!server_is_ready) { sleep(1); }

  struct sockaddr_in serv_addr = { 0 };
  serv_addr.sin_family = AF_INET;
  serv_addr.sin_port = htons(kServerPort);
  inet_aton(kServerIpAddress, &serv_addr.sin_addr);

  int sockfd = socket(AF_INET, SOCK_STREAM, 0);
  if (sockfd < 0)
    perror_and_exit("client socket");
  printf("[client] created socket: fd = %d\n", sockfd);

  printf("[client] connecting to server at %s:%d ...\n", kServerIpAddress, kServerPort);
  int retcode = connect(sockfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr));
  if (retcode < 0)
    perror_and_exit("connect");
  printf("[client] connected to server at %s:%d\n", kServerIpAddress, kServerPort);

  while (true) {
    char receive_buffer[1025] = { '\0' };
    printf("[client] calling read on socket\n");
    int bytes_read = read(sockfd, receive_buffer, sizeof(receive_buffer) - 1);
    if (bytes_read < 0) {
      perror_and_exit("read");
      exit(1);
    } else if (bytes_read == 0) {
      printf("[client] EOF\n");
      break;
    }
    receive_buffer[bytes_read] = '\0';
    printf("[client] received: '%s' (%d bytes)\n", receive_buffer, bytes_read);
  }

  printf("[client] DONE\n");
}

// Handle differing theread function signature.
#ifdef __barrelfish__
typedef int thread_return_t;
#else
typedef void* thread_return_t;
#endif

thread_return_t thread_adapter(void* arg) {
  typedef void (*runnable)(void);
  runnable func = (runnable) arg;
  func();
  return 0;
}

int main(int argc, char **argv) {
  if (argc != 2) {
    fprintf(stderr, "usage: %s <mode>\n", argv[0]);
    exit(1);
  }

  bool opt_server = false;
  bool opt_client = false;
  if (strcmp(argv[1], "server") == 0) {
    opt_server = true;
  } else if (strcmp(argv[1], "client") == 0) {
    opt_client = true;
    server_is_ready = true;  // not running in same process.
  } else if (strcmp(argv[1], "both") == 0) {
    opt_client = true;
    opt_server = true;
  } else {
    fprintf(stderr, "error: unknown argument '%s'\n", argv[1]);
    exit(1);
  }

#ifdef __barrelfish__
  struct thread* server_thread = NULL;
  struct thread* client_thread = NULL;
  if (opt_server)
    server_thread = thread_create(thread_adapter, run_server);
  if (opt_client)
    client_thread = thread_create(thread_adapter, run_client);

  if (server_thread != NULL)
    thread_join(server_thread, NULL);
  if (client_thread != NULL)
    thread_join(client_thread, NULL);
#else // !__barrelfish__
  pthread_t server_thread, client_thread;
  if (opt_server)
    pthread_create(&server_thread, NULL, thread_adapter, run_server);
  if (opt_client)
    pthread_create(&client_thread, NULL, thread_adapter, run_client);

  if (opt_server)
    pthread_join(server_thread, NULL);
  if (opt_client)
    pthread_join(client_thread, NULL);
#endif // __barrelfish__

  return 0;
}