/* By Hans Henrik Happe
 * 
 * usage: random-tcp <# processes> <# forwards>
 */

#include <asm/msr.h>

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

double second() {

  struct timeval tv;
  struct timezone tz;
  double t;

  gettimeofday(&tv,&tz);

  t= (double)(tv.tv_sec)+(double)(tv.tv_usec/1.0e6);

  return t;
}


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 do_read(int fd, void *buf, int n) {
    int res;
    
    res = read(fd, buf, n);
    if (res == -1) {
        perror("read");    
    }
    if (res != n) {
        fprintf(stderr, "read incomplete\n");    
    }
    
    return res;
}

int do_write(int fd, void *buf, int n) {
    int res;
    
    res = write(fd, buf, n);
    if (res == -1) {
        perror("write");    
    }
    if (res != n) {
        fprintf(stderr, "write incomplete\n");    
    }
    return res;
}


int main(int argc, char *argv[]) {
    int i, cnt, pid, dest, src, its;
    int lsock, sock;
    char id, rank, data;
    int port = 11100;
    double t0, t1;
        
    /* # processes */
    cnt = atoi(argv[1]);
    
    /* # forwards */
    its = atoi(argv[2]);
    
    {
        int socks[cnt];    
           
        /* Create processes */
        rank = 0;
        for (i=1; i<cnt; i++) {
            pid = fork();
            if (pid == 0) {
                rank=i;
                break;    
            }
        }

        /* Listen */
        lsock = start_listen(port+rank);

        
        sleep(2); /* "Ensure" that all processes are listening, HACK!!! */
                
        for (i=0; i<rank; i++) {
            sock = do_accept(lsock);
            
            do_read(sock, &id, 1);
            do_write(sock, &rank, 1);

            socks[id] = sock;
        }
        
        for (i=rank; i<cnt-1; i++) {
            sock = do_connect(port+i+1);
            
            do_write(sock, &rank, 1);
            do_read(sock, &id, 1);
            
            socks[id] = sock;
        }
                
        srandom(666);        
        src = 0;
        dest = 1;
        
        /* Write startup message */
        if (rank == src) {
            do_write(socks[dest], &data, 1);
        }
        
        
        /* Receive message and forward to a random destination */
        
        t0 = second();
        for (i=0; i < its; i++) {
                    
            if (rank == dest) {
                do_read(socks[src], &data, 1);
            }

            src = dest;                    
                    
            dest = random()%cnt;
                    
            /* Do not send to self */           
            if (dest == src) {
                dest = (dest+1)%cnt;    
            }
                    
            if (src == rank) {                   
               do_write(socks[dest], &data, 1);
            }
        }
        
        if (rank == 0) {
            for (i=1; i<cnt; i++) {
                wait(&dest);    
            }
            t1 = second();
            printf("avg forwarding time: %lf\n", ((t1-t0)/its));
        }
    }    
    
    
    return 0;
}