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

char buf[4096];
char tcp_buffer[4096];
int tcp_offset = -1;
int tcp_packet_size = 0;

int udp_socket(int port) {
    int s2;
    struct sockaddr_in si_me;

    if ((s2=socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP))==-1) {
        perror("socket");
        return -1;
    }
 
    // si_me stores our local endpoint. Remember that this program
    // has to be run in a network with UDP endpoint previously known
    // and directly accessible by all clients. In simpler terms, the
    // server cannot be behind a NAT.
    memset((char *) &si_me, 0, sizeof(si_me));
    si_me.sin_family = AF_INET;
    si_me.sin_port = htons(port);
    si_me.sin_addr.s_addr = htonl(INADDR_ANY);
    if (bind(s2, (struct sockaddr*)(&si_me), sizeof(si_me))==-1) {
        perror("bind");
        return -1;
    }
    return s2;
}

void init_address(struct sockaddr_in *sa, const char* host, int port) {
    memset((char *) sa, 0, sizeof(*sa));
    sa->sin_family = AF_INET;
    sa->sin_port = htons(port);
    if (inet_aton(host, &sa->sin_addr)==0) {
        perror("inet_aton");
    }
}

int tcp_socket(int port) {
    int s2;
    struct sockaddr_in si_me;

    if ((s2=socket(AF_INET, SOCK_STREAM, IPPROTO_TCP))==-1) {
        perror("socket");
        return -1;
    }
    
    int one=1;
    setsockopt(s2, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
 
    memset((char *) &si_me, 0, sizeof(si_me));
    si_me.sin_family = AF_INET;
    si_me.sin_port = htons(port);
    si_me.sin_addr.s_addr = htonl(INADDR_ANY);
    if (bind(s2, (struct sockaddr*)(&si_me), sizeof(si_me))==-1) {
        perror("bind");
        return -1;
    }
    
    listen(s2, 1);
    return s2;
}

int main(int argc, char* argv[])
{
    if (argc<3) {
        fprintf(stderr, "Usage: udppair_sip port port1 port2 ...\n");
        return 1;
    }
    int port1 = atoi(argv[1]);
    int ports[256];
    int sockets[256];
    struct sockaddr_in peers[256];
    int portcount = argc-2;
    int i;
    for(i=0; i<portcount; ++i) {
        char* arg = argv[i+2];
        if(strchr(arg, ':') == NULL) {
            ports[i]=atoi(arg);
            sockets[i]=udp_socket(ports[i]);
            fcntl(sockets[i], F_SETFL, O_NONBLOCK);
            memset((char *) &peers[i], 0, sizeof(peers[i]));
        } else {
            ports[i]=0;
            sockets[i]=udp_socket(0);
            fcntl(sockets[i], F_SETFL, O_NONBLOCK);
            char* sep = strchr(arg, ':');
            *sep='\0';
            init_address(&peers[i], arg, atoi(sep+1));
        }
    }

    int s1 = udp_socket(port1);
    int s1_t = tcp_socket(port1);
    int s1_tc=-1;

    fcntl(s1, F_SETFL, O_NONBLOCK);
    fcntl(s1_t, F_SETFL, O_NONBLOCK);

    fd_set rfds;

    struct sockaddr_in peer1;
    memset((char *) &peer1, 0, sizeof(peer1));

    for(;;) {
select_repeat:
        FD_ZERO(&rfds);
        FD_SET(s1, &rfds);
        FD_SET(s1_t, &rfds);
        int maxfd = s1_t;
        for(i=0; i<portcount; ++i) {
            FD_SET(sockets[i], &rfds);
            if (maxfd < sockets[i]) maxfd = sockets[i];
        }
        if(s1_tc != -1) {
            FD_SET(s1_tc, &rfds);
            if (maxfd < s1_tc) maxfd = s1_tc;
        }
        
        int ret = select(maxfd+1, &rfds, NULL, NULL, NULL);

        if(ret==-1) {
            if(errno==EINTR  || errno==EAGAIN) goto select_repeat;
            perror("select");
            return 2;
        }
        
        if(FD_ISSET(s1_t, &rfds)) {
            struct sockaddr_in a;
            socklen_t al = sizeof a;
            int q = accept(s1_t, (struct sockaddr*)&a, &al);
            if (q!=-1) {
                write(1, "t", 1);
                fcntl(q, F_SETFL, O_NONBLOCK);
                if(s1_tc!=-1) close(s1_tc);
                tcp_offset=-1;
                s1_tc = q;
            }
        }
        if(FD_ISSET(s1_tc, &rfds)) {
            if (tcp_offset==-1) {
                ret = recv(s1_tc, buf, 2, 0);
                if(ret!=2) {
                    close(s1_tc);
                    tcp_offset=-1;
                    s1_tc=-1;
                    write(1, "x", 1);
                    continue;
                }
                int packet_size = ntohs(*(unsigned short int*)buf);
                
                if (packet_size == 0xFFFF) { /* ping */
                    write(1, "p", 1);
                    ret = send(s1_tc, buf, 2, 0); 
                    if (ret != 2) {
                        close(s1_tc);
                        tcp_offset=-1;
                        s1_tc=-1;
                        write(1, "X", 1);
                    }
                } else {
                    tcp_packet_size = packet_size;
                    tcp_offset = 0;
                }
            } else {
                ret = recv(s1_tc, tcp_buffer+tcp_offset, tcp_packet_size-tcp_offset, 0);
                
                if(ret==0 || (ret == -1 && errno != EAGAIN && errno != EINTR)) {
                    close(s1_tc);
                    s1_tc=-1;
                    tcp_offset=-1;
                    write(1, "X", 1);
                    continue;
                }
                
                tcp_offset+=ret;
                
                if (tcp_offset == tcp_packet_size) {
                    tcp_offset = -1;
                    
                    socklen_t slen = sizeof(peer1);
                    write(1, "}", 1);
                    write(1, tcp_buffer, 1);
                    int i = (((unsigned char)tcp_buffer[0])-'0')%256;
                    ret = sendto(sockets[i], tcp_buffer+1, ret-1, 0, (struct sockaddr*)(&peers[i]), slen);
                    if (!ret) {
                        perror("sendto");
                    }
                }
            }
        }
        if(FD_ISSET(s1, &rfds)) {
            socklen_t slen = sizeof(peer1);
            ret = recvfrom(s1, buf, sizeof(buf), 0, (struct sockaddr*)(&peer1), &slen);
            if(ret==-1) {
                perror("recvfrom");
                return 4;
            }
            if(ret) {
                if (s1_tc != -1) {
                    write(1, "u", 1);
                    close(s1_tc);
                    s1_tc = -1;
                    tcp_offset=-1;
                }
                
                write(1, ">", 1);
                write(1, buf, 1);
                int i = (((unsigned char)buf[0])-'0')%256;
                ret = sendto(sockets[i], buf+1, ret-1, 0, (struct sockaddr*)(&peers[i]), slen);
                if (!ret) {
                    perror("sendto");
                } else {
                }
            }
        }
        for(i=0; i<portcount; ++i) {
        if(FD_ISSET(sockets[i], &rfds)) {
            socklen_t slen = sizeof(peer1);
            
            struct sockaddr_in tmpaddr;
            struct sockaddr* addrptr = (struct sockaddr*)&peers[i];
            if(!ports[i]) addrptr =  (struct sockaddr*)&tmpaddr; // don't allow override this address
                
            ret = recvfrom(sockets[i], buf+1, sizeof(buf)-1, 0, addrptr, &slen);
            if(ret==-1) {
                perror("recvfrom");
                return 4;
            }
            buf[0]=(i+'0')%256;
            
            if (s1_tc == -1) {
                ret = sendto(s1, buf, ret+1, 0, (struct sockaddr*)(&peer1), slen);
                if (!ret || ret == -1) {
                    perror("sendto");
                } else {
                    write(1, "<", 1);
                    write(1, buf, 1);
                }
            } else {
                int packet_length = ret+1;
                short int packet_len = htons(packet_length);
                ret = send(s1_tc, &packet_len, 2, 0);
                if (ret!=2) {
                    close(s1_tc);
                    s1_tc=-1;
                    tcp_offset=-1;
                    write(1, "!", 1);
                }
                ret = send(s1_tc, buf, packet_length, 0);
                if (ret != packet_length) {
                    close(s1_tc);
                    tcp_offset=-1;
                    s1_tc=-1;
                    write(1, "*", 1);
                } else {
                    write(1, "{", 1);
                    write(1, buf, 1);
                }
            }
        }
        }
    }    
}