#include <iostream>

#include <string.h> // for memset
#include <signal.h>

#include <errno.h>
#include <netdb.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#define MAX_SIZE 1024

/* write count bytes from buffer to fd, but handle
   partial writes (e.g. if the socket disconnects after
   some data is written)
 */
bool write_fully(int fd, const char *buffer, ssize_t count) {
    const char *ptr = buffer;
    const char *end = buffer + count;
    while (ptr != end) {
        ssize_t written = write(fd, (void*) ptr, end - ptr);
        if (written == -1) {
            return false;
        } 
        ptr += written;
    }
    return true;
}

void server_for_connection(int socket_fd) {
    int read_count;
    char request_buf[MAX_SIZE];
    while (1) {
        /* read as much as possible into request_buf */
        read_count = read(socket_fd, (void*) request_buf, MAX_SIZE);

        /* check for error, EOF */
        if (read_count == 0) {
            return; // EOF
        } else if (read_count == -1) {
            std::cerr << "read: " << strerror(errno) << std::endl;
            return;
        }

        std::cout << "read [" << std::string(request_buf, read_count) << "]\n";

        /* write it back; note that we don't care whether we read whole lines, etc. */
        if (false == write_fully(socket_fd, request_buf, read_count)) {
            std::cerr << "write: " << strerror(errno) << std::endl;
            return;
        }
    }
}

int make_server_socket(const char *hostname, const char *portname) {
    struct addrinfo *server;
    struct addrinfo hints;
    int rv;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_flags = AI_PASSIVE;
    rv = getaddrinfo(hostname, portname, &hints, &server);
    if (rv != 0) {
        std::cerr << "getaddrinfo: " << gai_strerror(rv) << std::endl;
        return -1;
    }
    int fd = -1;
    for (struct addrinfo *addr = server; addr; addr = addr->ai_next) {
        fd = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
        if (fd == -1)
            continue;
        /* make it so that we can bind() to the same address immediately
           after the server termiantes. By default, the address will be
           reserved for a while afterwards. */
        int one = 1;
        if(setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1)
            std::cerr << "setsockopt (SO_REUSEADDR): " << strerror(errno) << std::endl;
        if (bind(fd, addr->ai_addr, addr->ai_addrlen) == 0) {
            // XXX: really should handle getaddrinfo() returning multiple names
            break;
        }
        close(fd);
        fd = -1;
    }

    if (fd == -1) {
        std::cerr << "could not find address to bind to" << std::endl;
        return -1;
    }
    

    /* start listening for connections */
    listen(fd, SOMAXCONN);

    freeaddrinfo(server);

    return fd;
}

void output_server_host(int server_fd) {
    struct sockaddr_storage my_addr;
    socklen_t my_addr_len = sizeof(my_addr);
    getsockname(server_fd, (struct sockaddr*) &my_addr, &my_addr_len);
    char host[1024], port[1024];
    int rv = getnameinfo(
        (struct sockaddr*) &my_addr, my_addr_len,
        host, sizeof(host),
        port, sizeof(port),
        NI_NUMERICSERV | NI_NUMERICHOST
    );
    if (rv == 0) {
        std::cout << "Server bound to host " << host << ", port " << port << std::endl;
        std::cout << "(0.0.0.0 or :: represent all possible host addresses)" << std::endl;
    } else {
        std::cerr << "getnameinfo: " << gai_strerror(rv) << std::endl;
    }
}

int main(int argc, char **argv) {
    // ignore SIGPIPE, so writing to a socket the other end has closed causes
    // write() to return an error rather than crashing the program
    signal(SIGPIPE, SIG_IGN);
    const char *portname = NULL;
    const char *hostname = NULL;
    if (argc == 2) {
        portname = argv[1];
    } else if (argc == 3) {
        hostname = argv[1];
        portname = argv[2];
    } else {
        std::cerr << "usage: " << argv[0] << " PORT\n";
        std::cerr << "   or: " << argv[0] << " HOST PORT\n";
        std::cerr << "example: \n"
                  << argv[0] << " 127.0.0.1 9999\n";
        exit(1);
    }

    int server_fd = make_server_socket(hostname, portname);
    output_server_host(server_fd);
    if (server_fd != -1) {
        while (1) {
            struct sockaddr_storage remote_addr;
            socklen_t remote_addr_len = sizeof(remote_addr);
            int connection_fd = accept(server_fd, (struct sockaddr*) &remote_addr, &remote_addr_len);
            if (connection_fd == -1) {
                std::cerr << "accept: " << strerror(errno) << std::endl;
                exit(1);
            }
            char remote_host[1024], remote_port[1024];
            int rv = getnameinfo(
                (struct sockaddr*) &remote_addr, remote_addr_len,
                remote_host, sizeof(remote_host),
                remote_port, sizeof(remote_port),
                NI_NUMERICSERV | NI_NUMERICHOST
            );
            if (rv == 0) {
                std::cout << "Accepted connection from " << remote_host << ":" << remote_port << std::endl;
            } else {
                std::cerr << "getnameinfo: " << gai_strerror(rv);
            }
            server_for_connection(connection_fd);
            close(connection_fd);
            std::cout << "Closed connection, waiting for next" << std::endl;
        }
    }
}