#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <linux/un.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <errno.h>
#include <time.h>
#include <ctype.h>
#include <limits.h>

#define BUF_SIZE 1024

#undef IPPROTO_UDP
#define IPPROTO_UDP 0

#undef max
#define max(x,y) ((x) > (y) ? (x) : (y))

/*#define malloc(x) (aaa=malloc(x),v("m: %x %d %d", aaa, x, __LINE__),aaa)
#define free(x) (free(x),v("f: %x %d", x, __LINE__))
#define realloc(y,x) (aaa=realloc((y),(x)),v("re: %x %x %d %d", aaa, (y), (x), __LINE__),aaa)*/

int verbose = 0;

void *aaa;

char addr[64] = "224.0.0.213";
char db_fn[PATH_MAX] = "./db";
int port = 6969;
int hello_interval = 30;
int forget_cycles = 6;


int mcast_listen_fd = -1, mcast_read_fd = -1, mcast_write_fd = -1;
int local_listen_fd = -1, local_read_fd = -1, local_write_fd = -1;
char mcast_write_buf[BUF_SIZE], mcast_read_buf[BUF_SIZE];
char local_write_buf[BUF_SIZE], local_read_buf[BUF_SIZE];
int mcast_read_avail = 0, local_read_avail = 0;
int mcast_write_avail = 0, local_write_avail = 0;
int mcast_write_written = 0, local_write_written = 0;

time_t last_hello = -1;

void v(char *fmt, ...)
{
    if (verbose)
    {
        va_list ap;
        va_start(ap, fmt);
        vprintf(fmt, ap); 
        printf("\n");
        va_end(ap); 
    }
}
struct db_table
{
    char **entries;
    int n_entries;
    char *addr;
    int timeout;
};

struct s_db
{
    struct db_table **tables;
    int n_tables;
};

struct s_db db;


struct db_table *new_table(char *addr)
{
    struct db_table *table;

    table = malloc(sizeof (struct db_table));
    if (!table)
    {
        perror("malloc: table");
        exit(-1);
    }
    table->addr = addr;
    table->n_entries = 0;
    table->entries = 0;
    table->timeout = 0;

    db.tables = realloc(db.tables, (db.n_tables+1)*sizeof (struct db_table*));
    if (!db.tables)
    {
        perror("realloc: table");
        exit(-1);
    }
    db.tables[db.n_tables] = table;
    db.n_tables++;
    //printf("new table: %s %d %x %x\n", addr, db.n_tables, table, db.tables[db.n_tables-1]);

    return table;
}

void new_entry(struct db_table *table, char *entry)
{
    table->entries = realloc(table->entries, (table->n_entries+1)*sizeof(char *));
    if (!table->entries)
    {
        perror("realloc: entries");
        exit(-1);
    }
    table->entries[table->n_entries] = entry;
    table->n_entries++;
}

struct db_table *get_table(char *addr)
{
    int i;
    for (i=0; i<db.n_tables; i++)
    {
        if (!strcmp(addr, db.tables[i]->addr))
            return db.tables[i];
    }
    return 0;
}

void remove_table(int i)
{
    struct db_table *table;
    table = db.tables[i];
    if (i != db.n_tables-1)
    {
        db.tables[i] = db.tables[db.n_tables-1];
    }
    db.n_tables--;

    for (i=0; i<table->n_entries; i++)
    {
        char *entry = table->entries[i];
        //printf("!!! remove table: %d, %d\n", i, entry);
        if (entry)
            free(entry);
    }
    free(table->entries);
    free(table);
}

void remove_table_by_addr(char *addr)
{
    int i;
    for (i=0; i<db.n_tables; i++)
    {
        if (!strcmp(addr, db.tables[i]->addr))
            remove_table(i);
    }
}

void timeout_tables()
{
    int i;
    for (i=1; i<db.n_tables; i++)
    {
        db.tables[i]->timeout++;
        if (db.tables[i]->timeout >= forget_cycles)
        {
            remove_table(i);
            v("removed %s due to inactivity", db.tables[i]->addr);
        }
    }
    
}


static int recv_socket(int listen_port, char *mcast_address)
{
    struct sockaddr_in a;
    struct ip_mreq mreq;
    int s;
    int yes;

    if ((s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1) {
        perror("socket");
        return -1;
    }
    yes = 1;
    if (setsockopt(s, SOL_SOCKET, SO_REUSEADDR,
            (char *) &yes, sizeof(yes)) == -1) {
        perror("setsockopt");
        close(s);
        return -1;
    }
    memset(&a, 0, sizeof(a));
    a.sin_port = htons(listen_port);
    a.sin_family = AF_INET;
    if (bind(s, (struct sockaddr *) &a, sizeof(a)) == -1) {
        perror("bind");
        close(s);
        return -1;
    }

    //mcast_address = 0;
    if (mcast_address)
    {
        mreq.imr_multiaddr.s_addr = inet_addr(mcast_address);
        mreq.imr_interface.s_addr = htonl(INADDR_ANY);

        if (setsockopt(s, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof mreq) < 0)
            perror("recv IP_ADD_MEMBERSHIP");

    }
    if (mcast_address)
        v("listen: :%d (in group %s), %d", listen_port, mcast_address, s);
    else
        v("listen: :%d, %d", listen_port, s);
    return s;
}

static int connect_socket(int connect_port, char *address)
{
    struct sockaddr_in a;
    int s;

    if ((s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1) {
        perror("socket");
        close(s);
        return -1;
    }

    memset(&a, 0, sizeof(a));
    a.sin_port = htons(connect_port);
    a.sin_family = AF_INET;

    if (!inet_aton(address, (struct in_addr *) &a.sin_addr.s_addr)) {
        perror("bad IP address format");
        close(s);
        return -1;
    }

    if (connect(s, (struct sockaddr *) &a, sizeof(a)) == -1) {
        perror("connect()");
        shutdown(s, SHUT_RDWR);
        close(s);
        return -1;
    }
    v("connect: %s:%d", address, connect_port);
    return s;
}

void unix_recv_socket()
{
    char *fn = "/tmp/pwz";
    struct sockaddr_un recv_un;

    if ((local_listen_fd = socket(AF_UNIX, SOCK_DGRAM, 0)) == -1)
        perror("socket");

    memset(&recv_un, 0, sizeof recv_un);

    recv_un.sun_family = AF_UNIX;
    strncpy(recv_un.sun_path, fn, UNIX_PATH_MAX);
    unlink(recv_un.sun_path);

    if (bind(local_listen_fd, (struct sockaddr*)&recv_un, sizeof recv_un) < 0)
        perror("bind");

}

void unix_send_socket()
{
    char *fn = "/tmp/pwy";
    struct sockaddr_un send_un;

    if ((local_write_fd = socket(AF_UNIX, SOCK_DGRAM, 0)) == -1)
        perror("socket");

    memset(&send_un, 0, sizeof send_un);

    send_un.sun_family = AF_UNIX;
    strncpy(send_un.sun_path, fn, UNIX_PATH_MAX);

    if (connect(local_write_fd, (struct sockaddr*)&send_un, sizeof send_un) == -1) {
        perror("connect");
        shutdown(local_write_fd, SHUT_RDWR);
        close(local_write_fd);
    }

}

int tables_cmp(struct db_table **table1, struct db_table **table2)
{
    return strcmp((*table1)->addr, (*table2)->addr);
}
int str_cmp(char **str1, char **str2)
{
    //printf("%s %s\n", *str1, *str2);
    return strcmp(*str1, *str2);
}

void handle_local_msg()
{
    char msg[10240];
    int i,j;
    //printf("%d\n", db.n_tables);
    qsort(db.tables, db.n_tables, sizeof(struct db_table*), tables_cmp);

    for (i=0; i<db.n_tables; i++)
    {
        printf("%s %d\n", db.tables[i]->addr, db.tables[i]->timeout);
        qsort(db.tables[i]->entries, db.tables[i]->n_entries, sizeof(char *), str_cmp);
        for (j=0; j<db.tables[i]->n_entries; j++)
        {
            printf("   %s\n", db.tables[i]->entries[j]);
        }
    }

    if(send(local_write_fd, msg, strlen(msg)+1, 0) == -1)
        perror("send");
}

void handle_mcast_msg(char *from)
{
    int n, check_n=0, i, valid = 1, cur;
    n = ntohl(*(int*)mcast_read_buf);

    for (i=4;i<mcast_read_avail;i++)
    {
        char c = mcast_read_buf[i];
        if (c == 0)
            check_n++;
        else if (c < 32 || c > 127)
        {
            valid = 0;
        }    
    }
    if (n != check_n) valid =0;
    if (!valid)
    {
        v("got invalid msg from %s, ignoring", from);
        return;
    }

    struct db_table *table;

    table = get_table(from);
    if (!table)
    {
        table = new_table(from);
    }
    if (table && !n)
    {
        remove_table_by_addr(from);
        v("removed %s due to emptiness", db.tables[i]->addr);
        return;
    }

    /* remove old entries */
    for (i=0; i<table->n_entries; i++)
    {
        char *entry = table->entries[i];
        if (entry)
        {
            printf("[%x]\n", entry);
            free(entry);
        }
    }
    free(table->entries);

    /* insert new ones */
    for (i=0,cur=4;i<n;cur+=strlen(mcast_read_buf+cur)+1,i++)
    {
        char *entry;
        entry = malloc(sizeof(char*)*strlen(mcast_read_buf+cur));
        if (!entry)
        {
            perror("malloc: entry");
            exit(1);
        }
        strcpy(entry, mcast_read_buf+cur);
        new_entry(table, entry);
    }
    table->timeout = 0;

    v("handled message from %s (%d entries)", from, n);
    
    //write(2, mcast_read_buf, mcast_read_avail);
}

void send_mcast_msg()
{
    int n, len = 4, nn;
    struct db_table *local_table;

    char *msg;
    int msg_cur;

    local_table = get_table("local");

    for (n=0; n<local_table->n_entries; n++)
    {
        char *entry = local_table->entries[n];
        len += strlen(entry)+1;
    }

    msg = malloc(len);
    msg_cur = 0;

    nn = htonl(n);
    memcpy(msg, &nn, 4);
    msg_cur += 4;
    for (n=0; n<local_table->n_entries; n++)
    {
        char *entry = local_table->entries[n];
        memcpy(msg+msg_cur, entry, strlen(entry)+1);
        msg_cur += strlen(entry)+1;
    }

    if(send(mcast_write_fd, msg, len, 0) == -1)
        perror("send");

    v("sent mcast msg", last_hello, time(NULL));
    last_hello = time(NULL);
}

void interval()
{
    send_mcast_msg();
    timeout_tables();
}

#define SHUT_FD(fd) do {                                \
                     if (fd>= 0) {                 \
                         shutdown(fd, SHUT_RDWR);   \
                         close(fd);                 \
                         fd = -1;                   \
                         v("SHUT_FD: %d", fd);     \
                     }                               \
                 } while (0)


int main(int argc, char *argv[])
{
    int opt;
    db.tables = 0;

    while ((opt = getopt(argc, argv, "hva:i:p:f:d:")) != -1) {
        switch (opt) {
        case 'v':
            verbose = 1;
            break;
        case 'a':
            strncpy(addr, optarg, 64);
            break;
        case 'p':
            port = atoi(optarg);
            break;
        case 'i':
            hello_interval = atoi(optarg);
            break;
        case 'f':
            forget_cycles = atoi(optarg);
            break;
        case 'd':
            strncpy(db_fn, optarg, PATH_MAX);
            break;
        case 'h':
        default:
            fprintf(stderr, "Usage: %s [-h] [-v] [-a addr] [-p port] [-i interval] [-f forget_cycles]\n",
                    argv[0]);
            exit(1);
        } 
    }

    struct db_table* local_table;
    FILE *fp;
    char buf[1024];

    local_table = new_table("local");

    fp = fopen(db_fn, "r");
    if (!fp)
    {
        perror("fopen: db_fn");
        exit(1);
    }
    while (fgets(buf, 1024, fp))
    {
        char *entry;
        entry = malloc(sizeof(char)*strlen(buf));
        if (!entry)
        {
            perror("malloc: entry");
            exit(1);
        }
        strcpy(entry, buf);
        entry[strlen(buf)-1]=0;
        new_entry(local_table, entry);
    }
    v("read db: %s", db_fn);

    //sprintf(mcast_write_buf, "hi");
    //mcast_write_avail = 3;

    mcast_listen_fd = recv_socket(port, addr);
    mcast_write_fd = connect_socket(port, addr);
    unix_recv_socket();

    if (mcast_listen_fd == -1)
        exit(1);

    send_mcast_msg();

    for (;;) {
        int r, nfds = 0;
        fd_set rd, wr, er;

        FD_ZERO(&rd);
        FD_ZERO(&wr);
        FD_ZERO(&er);
        FD_SET(mcast_listen_fd, &rd);
        FD_SET(mcast_listen_fd, &er);
        nfds = max(nfds, mcast_listen_fd);
        FD_SET(local_listen_fd, &rd);
        nfds = max(nfds, local_listen_fd);
        if (mcast_read_fd > 0 && mcast_read_avail < BUF_SIZE)
        {
            FD_SET(mcast_read_fd, &rd);
            nfds = max(nfds, mcast_read_fd);
            v("FD_SET: mcast_read_fd");
        }
        if (local_read_fd > 0 && local_read_avail < BUF_SIZE)
        {
            FD_SET(local_read_fd, &rd);
            nfds = max(nfds, local_read_fd);
            v("FD_SET: local_read_fd");
        }
        if (mcast_write_fd > 0 && mcast_write_avail - mcast_write_written > 0)
        {
            FD_SET(mcast_write_fd, &wr);
            nfds = max(nfds, local_write_fd);
            v("FD_SET: mcast_write_fd");
        }
        if (local_write_fd > 0 && local_write_avail - local_write_written > 0)
        {
            FD_SET(local_write_fd, &wr);
            nfds = max(nfds, local_write_fd);
            v("FD_SET: local_write_fd");
        }

        struct timeval tv;
        tv.tv_sec = hello_interval;
        tv.tv_usec = 0;
        r = select(nfds + 1, &rd, &wr, &er, &tv);

        v("select: r=%d", r);
        if (r == -1 && errno == EINTR)
            continue;

        if (r == -1) {
            perror("select()");
            exit(EXIT_FAILURE);
        }

        if (FD_ISSET(mcast_listen_fd, &rd)) {
            unsigned int l;
            struct sockaddr_in client_address;
            l = sizeof client_address;

            r = recvfrom(mcast_listen_fd, mcast_read_buf, BUF_SIZE, 0, (struct sockaddr *) &client_address, &l);
            if (r == -1) {
                perror("recvfrom: mcast_listen_fd");
            } else {
                mcast_read_avail = r;
                v("message from %s",
                        inet_ntoa(client_address.sin_addr));
            }
            char *in_addr;
            in_addr = malloc(sizeof(char*)*strlen(inet_ntoa(client_address.sin_addr)));
            if (!in_addr)
            {
                perror("malloc: in_addr");
                exit(1);
            }
            strcpy(in_addr, inet_ntoa(client_address.sin_addr));
            handle_mcast_msg(in_addr);
        }
        if (FD_ISSET(local_listen_fd, &rd)) {
            unsigned int l;
            struct sockaddr_in client_address;
            l = sizeof client_address;

            r = recvfrom(local_listen_fd, local_read_buf, BUF_SIZE, 0, (struct sockaddr *) &client_address, &l);
            if (r == -1) {
                perror("recvfrom: local_listen_fd");
            } else {
                local_read_avail = r;
                v("local message from %s\n",
                        inet_ntoa(client_address.sin_addr));
            }
            unix_send_socket();
            handle_local_msg();
        }

        if (mcast_write_fd && FD_ISSET(mcast_write_fd, &wr)) {
            r = write(mcast_write_fd, mcast_write_buf + mcast_write_written,
                       mcast_write_avail - mcast_write_written);
            if (r < 1)
                SHUT_FD(mcast_write_fd);
            else
                mcast_write_written += r;
        }
        if (local_write_fd && FD_ISSET(local_write_fd, &wr)) {
            r = write(local_write_fd, local_write_buf + local_write_written,
                       local_write_avail - local_write_written);
            if (r < 1)
                SHUT_FD(local_write_fd);
            else
                local_write_written += r;
        }

        /* check if write data has caught read data */

        if (mcast_write_written == mcast_write_avail)
            mcast_write_written = mcast_write_avail = 0;
        if (local_write_written == local_write_avail)
            mcast_write_written = mcast_write_avail = 0;

        if (time(NULL) >= last_hello + hello_interval)
            interval();


    }
    exit(0);

}