#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <poll.h>
#include <stdlib.h>

#define SOCKETCOUNT (1 + 120 * 2)
#define BUFFERSIZE 100 * 1024
#define DNS_TIMEOUT 10

struct attr {
    char buffer[BUFFERSIZE];
    int start;
    int end;
};

void initialize(struct pollfd *sockets, struct attr *attrs, int serverSocket) {
    for (int i = 0; i < SOCKETCOUNT; i++) {
        sockets[i].fd = -1;
        sockets[i].events = 0;

        attrs[i].start = 0;
        attrs[i].end = 0;
    }

    sockets[0].fd = serverSocket;
    sockets[0].events = POLLIN;
}

int acceptNewConnection(struct pollfd *sockets, struct attr *attrs, int connectionCount, struct sockaddr_in destinationAddress, int serverSocket) {
    struct sockaddr_in incomingAddress;
    int incomingAddressLength = sizeof(incomingAddress);

    int incomingSocket = accept(serverSocket, (struct sockaddr *) &incomingAddress,
                                (socklen_t * ) & incomingAddressLength);

    if (-1 == incomingSocket) {
        perror("accept error");
        exit(EXIT_FAILURE);
    }

    if (connectionCount < 120) {
        int outcomingSocket = socket(AF_INET, SOCK_STREAM, 0);

        if (-1 == outcomingSocket) {
            perror("Cannot create outccoming socket");
            exit(EXIT_FAILURE);
        }

        int err = connect(outcomingSocket, (struct sockaddr *) &destinationAddress, sizeof(destinationAddress));
        if (-1 == err) {
            perror("Cannot connect to remote");
            exit(EXIT_FAILURE);
        }

        for (int j = 1; j < SOCKETCOUNT; j += 2) {
            if (-1 == sockets[j].fd) {

                sockets[j].fd = incomingSocket;
                sockets[j + 1].fd = outcomingSocket;

                sockets[j].events = POLLIN;
                sockets[j + 1].events = POLLIN;

                attrs[j].start = 0;
                attrs[j].end = 0;

                attrs[j + 1].start = 0;
                attrs[j + 1].end = 0;

                return 0;
            }
        }
    } else {
        close(incomingSocket);
    }

    return 1;
}

int readFromSocket(struct pollfd *sockets, struct attr *attrs, int i, int other) {
    if (BUFFERSIZE != attrs[i].end)
    {
        int received = read(sockets[i].fd, attrs[i].buffer + attrs[i].end, BUFFERSIZE - attrs[i].end);

        switch (received) {
            case -1:
                perror("read error");
                exit(EXIT_FAILURE);
            case 0:
                close(sockets[i].fd);
                sockets[i].fd = -1;
                sockets[i].events = 0;

                if (0 == attrs[i].end) {
                    close(sockets[other].fd);
                    sockets[other].fd = -1;
                    sockets[other].events = 0;

                    return 0;
                }
                break;
            default:
                attrs[i].end += received;

                if (BUFFERSIZE == attrs[i].end) {
                    sockets[i].events = sockets[i].events & (~POLLIN);
                }

                sockets[other].events = sockets[other].events | POLLOUT;
                break;
        }
    }

    return 1;
}

int writeToSocket(struct pollfd *sockets, struct attr *attrs, int i, int other) {
    if (0 != attrs[other].end) {
        int written = send(sockets[i].fd, attrs[other].buffer + attrs[other].start,
                           attrs[other].end - attrs[other].start, 0);

        switch (written) {
            case -1:
                perror("write error");
                exit(EXIT_FAILURE);
            case 0:
                close(sockets[i].fd);
                sockets[i].fd = -1;
                sockets[i].events = 0;

                if (0 == attrs[i].end) {
                    close(sockets[other].fd);
                    sockets[other].fd = -1;
                    sockets[other].events = 0;

                    return 0;
                }
                break;

            default:
                attrs[other].start += written;

                if (attrs[other].start == attrs[other].end) {
                    if (-1 == sockets[other].fd) {
                        close(sockets[i].fd);
                        sockets[i].fd = -1;
                        sockets[i].events = 0;

                        return 0;
                    } else {
                        attrs[other].start = 0;
                        attrs[other].end = 0;

                        sockets[i].events = sockets[i].events & (~POLLOUT);
                        sockets[other].events = sockets[other].events | POLLIN;
                    }
                }
        }
    } else if (-1 == sockets[other].fd)
    {
        close(sockets[i].fd);
        sockets[i].fd = -1;
        sockets[i].events = 0;

        return 0;
    }

    return 1;
}

void work(int serverSocket, struct sockaddr_in destinationAddress) {
    struct pollfd *sockets = (struct pollfd*) malloc(SOCKETCOUNT * sizeof(struct pollfd));
    struct attr *attrs = (struct attr*) malloc(SOCKETCOUNT * sizeof(struct attr));

    initialize(sockets, attrs, serverSocket);

    int connectionCount = 0;

    while(1) {
        int readyCount = poll(sockets, SOCKETCOUNT, -1);

        if (-1 == readyCount) {
            perror("poll error");
            exit(EXIT_FAILURE);
        }
        else if (0 == readyCount) {
            continue;
        }

        if (0 != (sockets[0].revents & POLLIN)) {
            if (0 == acceptNewConnection(sockets, attrs, connectionCount, destinationAddress, serverSocket)) {
                connectionCount++;
            }
            readyCount--;
        }

        for (int i = 1; i < SOCKETCOUNT && readyCount > 0; i++) {
            int other;

            if (1 == i % 2) {
                other = i + 1;
            }
            else {
                other = i - 1;
            }

            if (0 != (sockets[i].revents & POLLIN)) {
                if (readFromSocket(sockets, attrs, i, other) == 0) {
                    connectionCount--;
                }

                readyCount--;
            }
            else if (0 != (sockets[i].revents & POLLOUT)) {
                if (writeToSocket(sockets, attrs, i, other) == 0) {
                    connectionCount--;
                }

                readyCount--;
            } else if (0 != (sockets[i].revents & POLLHUP)) {
                exit(EXIT_SUCCESS);
            }
        }

        if (connectionCount >= 510) {
            sockets[0].fd = -1;
        } else {
            sockets[0].fd = serverSocket;
        }
    }

    free(sockets);
    free(attrs);
}

void sighandler(int signum) {
    fprintf(stderr, "Dns request timed out\n");
    exit(EXIT_FAILURE);
}

struct sockaddr_in getDestinationAddress(const char *hostName, int remotePort) {
    signal(SIGALRM, sighandler);
    alarm(DNS_TIMEOUT);
    struct hostent *hostInfo = gethostbyname(hostName);
    alarm(0);

    if (NULL == hostInfo) {
        fprintf(stderr, "Cannot get host by name\n");
        exit(EXIT_FAILURE);
    }

    struct sockaddr_in destinationAddress;

    destinationAddress.sin_family = AF_INET;
    destinationAddress.sin_port = htons(remotePort);
    memcpy(&destinationAddress.sin_addr, hostInfo->h_addr, hostInfo->h_length);

    return destinationAddress;
}

int createServerSocket(int serverPort) {
    int serverSocket = socket(AF_INET, SOCK_STREAM, 0);

    if (-1 == serverSocket) {
        perror("Cannot create serverSocket");
        exit(EXIT_FAILURE);
    }

    struct sockaddr_in serverAddress;

    serverAddress.sin_family = AF_INET;
    serverAddress.sin_addr.s_addr = htonl(INADDR_ANY);
    serverAddress.sin_port = htons(serverPort);

    if (-1 == bind(serverSocket, (struct sockaddr *) &serverAddress, sizeof(serverAddress))) {
        perror("binding error");
        exit(EXIT_FAILURE);
    }

    if (-1 == listen(serverSocket, 1024)) {
        perror("listen error");
        exit(EXIT_FAILURE);
    }

    return serverSocket;
}

void checkCountArguments(int argc) {
    if (argc != 4) {
        perror("Wrong count of arguments");
        exit(EXIT_FAILURE);
    }
}

int main(int argc, char *argv[]) {
    checkCountArguments(argc);

    int serverPort = atoi(argv[1]);
    int remotePort = atoi(argv[3]);

    if (0 == serverPort) {
        perror("Incorrect port for listening");
        exit(EXIT_FAILURE);
    }

    struct sockaddr_in destinationAddress = getDestinationAddress(argv[2], remotePort);
    int serverSocket = createServerSocket(serverPort);

    work(serverSocket, destinationAddress);


    putchar('\n');

    return EXIT_SUCCESS;
}