#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <arpa/inet.h>
#include <sys/epoll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <unistd.h>

#ifndef PR_SET_TIMERSLACK
#define PR_SET_TIMERSLACK 29
#endif

const char *g_mcastaddr = "239.100.0.99";
int g_port = 10100;

void error(const char *s)
{
    fprintf(stderr, "%s\n", s);
    exit(1);
}

void check(int v)
{
    int myerr = errno;
    char *myerrstr = strerror(myerr);
    if(!v)
        error("bad return code");
}

int main(int argc, char **argv)
{
    if(argc != 2)
        error("usage: mcasttest (server|client)");
    if(strcmp(argv[1], "client") == 0)
    {
        /*
         * Client program: subscribes to a multicast group, receives messages
         * and prints a count of messages received once it's done.
         */
        int s = socket(AF_INET, SOCK_DGRAM, 0);
        check(s > 0);
        int val = 1;
        check(setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val)) == 0);

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_port = htons(g_port);
        addr.sin_addr.s_addr = htonl(INADDR_ANY);
        check(bind(s, (struct sockaddr *) &addr, sizeof(addr)) == 0);

        struct ip_mreqn mreq;
        memset(&mreq, 0, sizeof(mreq));
        check(inet_pton(AF_INET, g_mcastaddr, &mreq.imr_multiaddr));
        mreq.imr_address.s_addr = htonl(INADDR_ANY);
        mreq.imr_ifindex = 0;
        check(setsockopt(s, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) == 0);

        int bufSz;
        socklen_t len = sizeof(bufSz);
        getsockopt(s, SOL_SOCKET, SO_RCVBUF, (char*)(&bufSz), &len);
        printf("bufsz: %d\n", bufSz);

        int npackets = 0;
        char buf[1000];
        memset(buf, 0, sizeof(buf));
        while(1)
        {
            struct sockaddr_in from;
            socklen_t fromlen = sizeof(from);
            check(recvfrom(s, buf, 1000, 0, (struct sockaddr*)&from, &fromlen) == 100);
            ++npackets;
            if(buf[0] == 1) // exit message
                break;
        }
        printf("received %d packets\n", npackets);
    }
    else if(strcmp(argv[1], "server") == 0)
    {
        /*
         * Setup a timer resolution of 1000 ns : 1 us
         */
        prctl(PR_SET_TIMERSLACK, 1000); 

        /*
         * Server program: sends 50,000 packets per second to a multicast address,
         * for 10 seconds.
         */
        int s = socket(AF_INET, SOCK_DGRAM, 0);
        int val = 1;
        int i = 1;
        check(s > 0);

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_port = htons(g_port);
        check(inet_pton(AF_INET, g_mcastaddr, &addr.sin_addr.s_addr));
        check(connect(s, (struct sockaddr *) &addr, sizeof(addr)) == 0);

        int npackets = 500000;
        char buf[100];
        memset(buf, 0, sizeof(buf));
        for(i = 1; i < npackets; ++i)
        {
            check(send(s, buf, sizeof(buf), 0) > 0);
            usleep(20); // 50,000 messages per second
        }

        buf[0] = 1;
        for(i = 1; i < 5; ++i)
        {
            check(send(s, buf, sizeof(buf), 0) > 0);
            sleep(1);
        }
    }
    else
        error("unknown mode");
    return 0;
}