/* 
 * usage: random-inet <# processes> <# messages>
 */


#include <stdio.h>
#include <poll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/tcp.h>
#include <fcntl.h>
#include <netdb.h>


int do_connect(int port) {
   int n, sock, on=1;
   struct addrinfo hints, *res;
   char str[6];
   void *adr;
    
   memset(&hints, 0, sizeof(struct addrinfo));
    
   hints.ai_flags    = AI_PASSIVE;
   hints.ai_family   = PF_UNSPEC;
   hints.ai_socktype = SOCK_STREAM;

   sprintf(str, "%d", port);
   n = getaddrinfo("localhost", str, &hints, &res);

   if (n != 0) {
       fprintf(stderr,
               "getaddrinfo error: [%s]\n",
               gai_strerror(n));
       return -1;    
   }
   
   sock = socket(AF_INET, SOCK_STREAM, 0);
   if (sock == -1) {
       perror("socket");
       return -1;
   }

    
   if (setsockopt(sock, SOL_TCP, TCP_NODELAY, &on, sizeof(on)) == -1) {
       perror("setsockopt");
       return -1;
   }
   
   if (connect(sock, (struct sockaddr *)res->ai_addr, sizeof(*res->ai_addr)) == -1) {
       perror("connect");
       return -1;
   }
   
   freeaddrinfo(res);

   return sock;
}

int start_listen(int port) {
    int n, on=1;
    int sock;    
    struct sockaddr_in name;
   
        
    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == -1) {
        perror("socket");
        return -1;
    }

    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) == -1) {
        perror("setsockopt");
        return -1;
    }
        
    name.sin_family = AF_INET;
    name.sin_port = htons (port);
    name.sin_addr.s_addr = htonl (INADDR_ANY);
    
    if (bind (sock, (struct sockaddr *) &name, sizeof (name)) == -1) {
        perror("bind");
        return -1;
    }      
    
    if (listen(sock, 10) == -1) {
        perror("listen");
        return -1;
    }
    
    return sock;  
}
        
int do_accept(int lsock) {
    struct sockaddr addr;
    socklen_t len = sizeof(addr);
    int sock, on=1;


    if ((sock = accept(lsock, &addr, &len)) == -1) {
        perror("accept");
        return -1;
    }

    if (setsockopt(sock, SOL_TCP, TCP_NODELAY, &on, sizeof(on)) == -1) {
        perror("setsockopt");
        return -1;
    }
        
    return sock;
}
        
                


int main(int argc, char *argv[]) {
    int i, j, n, cnt, pid, rank, dest;
    int lsock;
    char data = 'h';
    int port = 11100;
    
    /* # processes */
    cnt = atoi(argv[1]);
    
    /* # messages */
    n = atoi(argv[2]);

    {
        int socks[cnt-1];    
        struct pollfd pfds[cnt-1];
           
        /* Create processes */
        rank = 0;
        for (i=1; i<cnt; i++) {
            pid = fork();
            if (pid == 0) {
                rank=cnt-i;
                break;    
            }
        }

        /* Setup connections */
        
        lsock = start_listen(port+rank);
        
        for (i=0; i<rank; i++) {
            socks[i] = do_accept(lsock);
            pfds[i].fd = socks[i];
            pfds[i].events = POLLIN;
        }
        
        for (i=rank; i<cnt-1; i++) {
            socks[i] = do_connect(port+i+1);
            pfds[i].fd = socks[i];
            pfds[i].events = POLLIN;
        }
                
        srandom(rank);        
        
        /* Write startup messages */
        if (rank < n) {
            dest = random()%(cnt-1);
            write(socks[dest], &data, 1);
        }

        /* Receive and forward messages to random destinations */
        while (1) {
            poll(pfds, cnt-1, -1);
            j = 0;
            for (i=0; i<cnt-1; i++) {
                if (pfds[i].revents != 0) {
                    read(pfds[i].fd, &data, 1);
                    j++;
                }
            }
            
            for (i=0; i<j; i++) {
                dest = random()%(cnt-1);
                write(socks[dest], &data, 1);
            }
        }
    }    
    return 0;
}