#include <unistd.h>
#include <sys/wait.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <time.h>
#include "signal_action.h"
#include "file_lock.h"

#include "babysit.h"

/* TODO:
 * - instead of using bytes, maybe use lines of text. But bytes are
 *   easier right now
 */

/* Thanks to Dean Gaudet for the idea to use a fd to replace a signal */
int sigchld_pipe[2];
#define READ_END 0
#define WRITE_END 1

/* child_pid == 0 means no child is running */
pid_t child_pid = 0;

/* Kill if pid != 0 */
#define kill_if(a, b) ((a) ? kill(a, b) : 0)

/* takes an fd.
 * Returns it if it can be close-on-execed. 
 * Otherwise, closes fd and returns -1
 */
int coe(int fd) {
    if (fd == -1) {
        return fd;
    }
    if (fcntl(fd, F_SETFD, FD_CLOEXEC) == 0) {
        return fd;
    }
    close(fd);
    return -1;
}

/* Check if the current working directory is valid.
 * return 1 on success, 0 on failure, sets errno
 */
int valid_wd(void) {
    struct stat wd;

    if (stat(".", &wd)) {
        return 0;
    }
    if (wd.st_nlink == 0) {
        errno = ENOENT;
        return 0;
    }
    return 1;
}
        
void sigchld_handler(int sig) {
    /* Have to worry about pids other than child_pid, because of a shell
     * script like:
     *
     * xterm &
     * exec babysit /blah/blah
     *
     * followed by an exit of the xterm. Have to waitpid(-1) to avoid
     * zombies.
     */

    const char *byte_ptr = "";
    pid_t pid;
       
    /* Looping here because there isn't necessarily one signal per child
     * in the queue */
    while ((pid = waitpid(-1, NULL, WNOHANG)) > 0) {
        if (pid == child_pid) {
            write(sigchld_pipe[WRITE_END], byte_ptr, 1);
            child_pid = 0;
        }
    }
    if ((pid == -1) && (errno != ECHILD)) {
        perror("waitpid");
        exit(1);
    }
}

void child(void) {
    char *const run_argv[] = {"./run", NULL};

    execv(*run_argv, run_argv);
    /* Must have failed */
    perror("execv ./run");
    exit(1);
}

/* Sleep at least as long as we tell it to, even through signals */
void sleep_dammit(int sec, int nsec) {
    struct timespec want; /* Time we want to sleep */
    struct timespec remain; /* Time remaining to sleep */

    want.tv_sec = sec;
    want.tv_nsec = nsec;
    while (nanosleep(&want, &remain) != 0) {
        if (errno == EINTR) {
            want = remain;
            continue;
        }
        perror("nanosleep");
        exit(1);
    }
}

/* Make a pipe that disappears on exec */
int no_exec_pipe(int filedes[2]) {
    int rc;

    rc = pipe(filedes);
    if (rc != 0) return rc;
    if ((coe(filedes[0]) == -1) || (coe(filedes[1]) == -1)) {
        return -1;
    }
    return 0;
}

/* Make a lock (+ file) that disappears on exec */
int no_exec_file_trylock(const char *lock_file_name) {
    int fd;

    fd = coe(open(lock_file_name, O_WRONLY|O_CREAT, 0600));
    if (fd == -1) {
        perror("open lock file");
        exit(1);
    }

    if (file_trylock(fd) != 0) {
        close(fd);
        return -1;
    }

    return fd;
}

int main(int argc, const char * const argv[]) {
    /* argv[1] is always set, so this is safe */
    const char *service_dir = argv[1];
    int fifo_fd;
    int rc;
    struct stat unused_stat_struct;
    char byte;
    struct pollfd pollarray[2];
/* For pollarray */
#define CONTROL_FIFO_SLOT 0
#define SIGCHLD_PIPE_SLOT 1
    int want_exit = 0;          /* Exit when child exits */
    int want_new_child = 1;     /* Start new child when child exits */
    int keep_wanting_new_child = 1; /* Maintain want_new_child when child exits */
    int need_a_break = 0;       /* Sleep 1 second after fork */

    sigset_t sigfullmask, sigemptymask;
    sigfillset(&sigfullmask);
    sigemptyset(&sigemptymask);

    /* Blocking signals everywhere possible */
    sigprocmask(SIG_SETMASK, &sigfullmask, NULL);

    if (argc != 2) {
       fprintf(stderr, "Usage: %s [service_dir]\n", argv[0]);
       exit(1);
    }

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

    if (chdir(service_dir) != 0) {
        fprintf(stderr, "unable to chdir to %s: %s\n", service_dir, strerror(errno));
        exit(1);
    }

    if ((mkdir(BABYSIT_DIR, 0700) != 0) && (errno != EEXIST)) {
        perror("mkdir");
        exit(1);
    }

    file_lock_init();
    /* Don't need the fd, the unlock will happen on exit */
    if (no_exec_file_trylock(LOCK_FILE) == -1) {
        perror("Can't lock file, babysit is probably running already");
        exit(1);
    }

    if ((mkfifo(CONTROL_FIFO, 0600) != 0) && (errno != EEXIST)) {
        perror("mkfifo control fifo");
        exit(1);
    }

    /* Need to open both a read and write fd to the pipe so that it
     * doesn't go EOF on us. Linux lets us accomplish this with an
     * O_RDWR pipe, but others don't. We don't need the writable fd */
    fifo_fd = coe(open(CONTROL_FIFO, O_RDONLY|O_NONBLOCK));
    if ((fifo_fd == -1) || (coe(open(CONTROL_FIFO, O_WRONLY)) == -1)) {
        perror("open control fifo");
        exit(1);
    }

    if (stat(DOWN_FILE, &unused_stat_struct) == 0) {
        want_new_child = 0;
    }
     
    if (no_exec_pipe(sigchld_pipe) != 0) {
        perror("no_exec_pipe");
        exit(1);
    }

    pollarray[SIGCHLD_PIPE_SLOT].fd = sigchld_pipe[READ_END];
    pollarray[SIGCHLD_PIPE_SLOT].events = POLLIN;
    pollarray[CONTROL_FIFO_SLOT].fd = fifo_fd;
    pollarray[CONTROL_FIFO_SLOT].events = POLLIN;

    while (1) {
        if (!child_pid) {
            if (want_exit) {
                exit(0);
            }
            /* If the cwd doesn't exist, we really can't do anything.
             * babysit/control is inaccessible, and so is the run
             * script, so just give up.
             * Users: Don't depend on this behavior. */
            if (!valid_wd()) {
                perror("Bad babysit directory");
                exit(1);
            }
            if (want_new_child) {
                child_pid = fork();
                if (child_pid == 0) {
                    sigprocmask(SIG_SETMASK, &sigemptymask, NULL);
                    child();
                }
                else if (child_pid < 0) {
                    perror("fork");
                    exit(1);
                }
                want_new_child = keep_wanting_new_child;
                need_a_break = 1;
            }
        }

        sigprocmask(SIG_SETMASK, &sigemptymask, NULL);
        if (need_a_break) {
            /* Wait one second after fork to avoid insanity */
            sleep_dammit(1, 0);
            need_a_break = 0;
        }
        /* Block until either a signal or a byte on the control fifo */
        do {
            rc = poll(pollarray, 2, -1);
        } while ((rc == -1) && (errno == EINTR));
        sigprocmask(SIG_SETMASK, &sigfullmask, NULL);

        if (rc == -1) {
            perror("poll");
            exit(1);
        }
        if ((pollarray[SIGCHLD_PIPE_SLOT].revents
             | pollarray[CONTROL_FIFO_SLOT].revents)
            & (POLLERR|POLLHUP|POLLNVAL)) {
            fputs("poll error\n", stderr);
            exit(1);
        }

        if (pollarray[SIGCHLD_PIPE_SLOT].revents & POLLIN) {
            /* Just clear out the pipe */
            if (read(pollarray[SIGCHLD_PIPE_SLOT].fd, &byte, 1) != 1) {
                fputs("sigchld pipe read error\n", stderr);
                exit(1);
            }
        }
        if (pollarray[CONTROL_FIFO_SLOT].revents & POLLIN) {
            if (read(pollarray[CONTROL_FIFO_SLOT].fd, &byte, 1) == 1) {
                /* All the known error conditions for kill are not
                 * useful to check for */
                switch (byte) {
                    case 'u':
                        want_new_child = keep_wanting_new_child = 1;
                        break;
                    case 'd':
                        want_new_child = 0;
                        kill_if(child_pid, SIGTERM);
                        kill_if(child_pid, SIGCONT);
                        break;
                    case 'o':
                        want_new_child = 0;
                        break;
                    case 'p':
                        kill_if(child_pid, SIGSTOP);
                        break;
                    case 'c':
                        kill_if(child_pid, SIGCONT);
                        break;
                    case 'h':
                        kill_if(child_pid, SIGHUP);
                        break;
                    case 'a':
                        kill_if(child_pid, SIGALRM);
                        break;
                    case 'i':
                        kill_if(child_pid, SIGINT);
                        break;
                    case 't':
                        kill_if(child_pid, SIGTERM);
                        break;
                    case 'k':
                        kill_if(child_pid, SIGKILL);
                        break;
                    case 'x':
                        want_exit = 1;
                        break;
                }
            }
            else {
                fputs("control fifo read error\n", stderr);
                exit(1);
            }
        }
    }
}