#include <pthread.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h> /* for setting the registration file descripter blocking */
#include <sys/ioctl.h> /* for ioctl() */
#include <sys/un.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <limits.h>
#include <errno.h>

#define AST(a) do { if (!(a)) {printf("error %d\n", __LINE__); return (void *)1;} } while(0)
#define MAST(a) do { if (!(a)) {printf("error %d\n", __LINE__); return 0;} } while(0)
#define MAX_CONNECT_TRIES 3
static char *PORT_ADDR = "check_sockets_port";
static int open_port(const char *address);
static int accept_a_connection(int register_socket, int *new_socket);
static int connect_to(const char *host);
static pthread_attr_t attr;
static int FORKED=0;
static int WAIT_ON_SELECT=1;
static int NON_BLOCKING_ACCEPT=1;
static int IOCTL=0;

void *server(void *addr_buf)
{
  int port_sock, cli_sock, len;
  port_sock = open_port(addr_buf);
  AST(port_sock >= 0);
  AST(accept_a_connection(port_sock, &cli_sock) == 0);
  AST(shutdown(cli_sock, SHUT_RDWR)==0);
  AST(close(cli_sock)==0);
  AST(close(port_sock)==0);
  AST(unlink(addr_buf)==0);
  return NULL;
}

void *client(void *addr_buf)
{
  int serv_sock, len;
  serv_sock = connect_to(addr_buf);
  AST(serv_sock >= 0);
  AST(shutdown(serv_sock, SHUT_RDWR)==0);
  AST(close(serv_sock)==0);
  return NULL;
}

int main (int argc, char **argv)
{
  int pid;
  pthread_t s_id, c_id;
  int s_ans, c_ans;
  if (argc > 1 && strlen(argv[1]) > 0) {
    if (strncmp(argv[1], "h", 2)==0) {
      printf("options: 'i' use ioctl instead of fcntl to set socket blocking\n");
      return 0;
    }
    if (strncmp(argv[1], "i", 2)==0)
      IOCTL = 1;
  }
  if (FORKED) {
    printf("Forked\n");
    pid = fork();
    if (pid > 0) { /* parent */ 
      server(PORT_ADDR);
    } else if (pid == 0) { /* child */
      sleep(1);
      client(PORT_ADDR);
    } else {
      abort();
    }
  } else {
    printf("Threaded\n");
    MAST(pthread_attr_init(&attr)==0);
    MAST(pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE)==0);
    MAST(pthread_create(&s_id, &attr, &server, PORT_ADDR)==0);
    sleep(1);
    MAST(pthread_create(&c_id, &attr, &client, PORT_ADDR)==0);
    MAST(pthread_join(s_id, (void **)&s_ans)==0);
    MAST(pthread_join(c_id, (void **)&c_ans)==0);
    MAST(s_ans==0);
    MAST(c_ans==0);
    MAST(pthread_attr_destroy(&attr)==0);
  }
  printf("Success\n");
  return 0;
}

/* create a new Unix domain socket and bind it to SERVER_HOST_ADDR */
static int open_port(const char *address)
{
  struct sockaddr_un    serv_addr = {0};
  int                   sock;

  /* Open a new port. Try to get requested port. */
  if ((sock = socket(AF_UNIX, SOCK_STREAM, 0)) < 0) {
    perror("socket");
    abort();
  }
  serv_addr.sun_family = AF_UNIX;
  strcpy( serv_addr.sun_path, address);
  unlink( address );
  if (bind(sock, (struct sockaddr *)&serv_addr, SUN_LEN(&serv_addr)) <0 ) {
    perror("bind");
    abort();
  }
  if (listen(sock, 5) < 0) {
    perror("listen");
    abort();
  }
  printf("Bound and listening to \"%s\" on fd %d\n", address, sock);
  fflush(stdout);
  return sock;
}

static void print_bits(int num, int len)
{
  int i, out;
  for (i=0; i<len; i++) {
    if (num & 1<<i)
      printf("%d", 1);
    else
      printf("%d", 0);
  }
  printf("\n");
}

static int
set_nonblock_flag (int desc, int nonblocking)
{
  int oldflags, flag;
  printf("set %d nonblock flag to %d with", desc, nonblocking);
  if (IOCTL) {
    printf(" ioctl()\n");
    return ioctl(desc, FIONBIO, &nonblocking); /* 1 for nonblocking */
  } else {
    printf(" fcntl()\n");
    oldflags = fcntl (desc, F_GETFL, 0);
    if (1) {
      printf("O_NONBLOCK\t");
      print_bits(O_NONBLOCK, 32);
      printf("oldflags  \t");
      print_bits(oldflags, 32);
    }
    /* If reading the flags failed, return error indication now. */
    if (oldflags == -1)
      return -1;
    /* Set just the flag we want to set. */
    if (nonblocking)
      oldflags |= O_NONBLOCK; /* set nonblocking */
    else
      oldflags &= ~O_NONBLOCK; /* set blocking */
    if (1) {
      printf("newflags  \t");
      print_bits(oldflags, 32);
    }
    /* Store modified flag word in the descriptor. */
    return fcntl (desc, F_SETFL, oldflags);
  }
}

/* not thread safe to block on accept in Cygwin, (see test ccs.exe)
   so we'll block on select instead */
static int accept_a_connection(int register_socket, int *new_socket)
{
  /* even with this select, still blocks on accept/connect */
  int num_ready;
  fd_set readset;
  if (WAIT_ON_SELECT) {
    printf("Wait on select\n");
    fflush(stdout);
    FD_ZERO(&readset);
    FD_SET(register_socket, &readset);
    num_ready = select(register_socket+1, &readset, 0, 0, 0);
    if (num_ready <= 0) perror("Select failed."), abort();
    if (!FD_ISSET(register_socket, &readset)) {
      abort();
    }
  } else {
    printf("Dont' wait on select\n");
  }
  if (NON_BLOCKING_ACCEPT) {
    if (set_nonblock_flag(register_socket, 1) == -1)
      perror("nonblock failed."), abort();
    printf("Try accept from %d\n", register_socket);
    fflush(stdout);
    *new_socket = accept(register_socket, NULL, NULL);
    if (*new_socket < 0)
      perror("accept failed."), abort();
    if (set_nonblock_flag(register_socket, 0) == -1)
      perror("block failed."), abort();
  } else {
    printf("Wait on accept from %d\n", register_socket);
    fflush(stdout);
    *new_socket = accept(register_socket, NULL, NULL);
  }
  if (*new_socket < 0) perror("Accept failed."), abort();
  printf("Accepted new connection on %d\n", *new_socket);
  fflush(stdout);
  return 0;
}


/* open a socket connecting to HOST, used for test sockets */
static int connect_to(const char *host)
{
  struct sockaddr_un     serv_addr = {0};
  int                    sock, j=0, err;
  
  while (1) { /* try a few times */
    if ((sock = socket(AF_UNIX, SOCK_STREAM, 0)) < 0) {
      perror("connect_to: can't open stream socket ");
      abort();
    }
    memset((char *) &serv_addr, 0, sizeof(serv_addr));
    serv_addr.sun_family = AF_UNIX;
    strcpy(serv_addr.sun_path, host);
    printf("Connecting %d to \"%s\"\n", sock, serv_addr.sun_path);
    fflush(stdout);
    if (connect(sock, (struct sockaddr *)&serv_addr,SUN_LEN(&serv_addr))>=0){
      printf("Connected :)\n");
      fflush(stdout);
      break; /* we got it :) */
    }
    printf("Connection failed\n");
    fflush(stdout);
    close(sock);
    sleep(1);	/* Let the server come up. */
  }
  return sock;
}