#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <sys/wait.h>
#include <errno.h>
#include "signal_action.h"

/* TODO
 * - should check for close() errors, theoretically
 * - what kind of error checking is needed for inet_aton, etc.?
 * - More debugging output
 */

int connections_left = 40;

static void sigchld_handler(int sig) {
    pid_t pid;

    /* Looping here because there isn't necessarily one signal per child
     * in the queue */
    while ((pid = waitpid(-1, NULL, WNOHANG)) > 0) {
        connections_left++;
    }
    if ((pid == -1) && (errno != ECHILD)) {
        perror("waitpid");
        exit(1);
    }
}

void child(int socket, char *const argv[]) {
    setsid();
    if ((dup2(socket, STDIN_FILENO) == -1)
        || (dup2(socket, STDOUT_FILENO) == -1)) {
        perror("dup");
        exit(1);
    }
    close(socket);
    execvp(*argv, argv);
    perror("execvp");
    exit(1);
}

int main(int argc, char *const argv[]) {
    int socket;
    pid_t pid;
    char *const *new_argv;
    int option;
    sigset_t sigfullmask, sigemptymask;
    sigfillset(&sigfullmask);
    sigemptyset(&sigemptymask);

    /* Blocking signals everywhere but in accept() and sigsuspend() */
    sigprocmask(SIG_SETMASK, &sigfullmask, NULL);

    if (signal_action(SIGCHLD, sigchld_handler) != 0) {
        perror("signal_action");
        exit(1);
    }

    while ((option = getopt(argc, argv, "c:")) != -1) {
        switch (option) {
            case 'c':
                connections_left = atoi(optarg);
                break;
        }
    }

    new_argv = argv + optind;
    if (*new_argv == NULL) {
        fputs("Must provide a command to execute\n", stderr);
        exit(1);
    }
    while (1) {
        while (connections_left < 1) {
            /* Wait until one of the children dies and we receive a SIGCHLD. */
            sigsuspend(&sigemptymask);
        }

        /* This should be the only place we block (other than sigsuspend
         * above). A good place to catch signals, if such a place exists */
        sigprocmask(SIG_SETMASK, &sigemptymask, NULL);
        do {
            socket = accept(STDIN_FILENO, NULL, NULL);
        } while ((socket < 0) && (errno == EINTR));
        sigprocmask(SIG_SETMASK, &sigfullmask, NULL);

        if (socket < 0) {
            perror("accept");
            exit(1);
        }
        pid = fork();
        if (pid == 0) {
            child(socket, new_argv);
        }
        else if (pid < 0) {
            perror("fork");
            exit(1);
        }
        close(socket);
        connections_left--;
    }
}