#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <syslog.h>
#include <errno.h>
#include <string.h>
#include <limits.h>
#include "ucore/supervisor.h"


#define CLEAN_EXIT 11
#define BAD_EXIT   22
#define RESTART    33

//seconds to wait before restarting the process
#define RESTART_GRACE 1

static volatile sig_atomic_t restart_flag = 0;
static volatile sig_atomic_t exit_flag = 0;

static void restart_handler(int sig)
{
    (void)sig;
    restart_flag = 1;
}

static void exit_handler(int sig)
{
    (void)sig;
    exit_flag = 1;
}

static void kill_supervised_process(pid_t pid)
{
    kill(pid, SIGTERM);
    /* Todo.. verify we killed it, otherwise take more 
     * desperate measures*/
}

static int watch(pid_t pid)
{
    int status;
    pid_t rc;

again:
    rc = waitpid(pid, &status, 0);
    if (rc == -1 && errno == EINTR) {
        if (restart_flag) {
            restart_flag = 0;
            kill_supervised_process(pid);
            /* next waitpid should notify us of the death and
             * we'll restart the process */
            goto again;
        } else if (exit_flag) {
            exit_flag = 0;
            return CLEAN_EXIT;
        }

        goto again;

    } else if (rc == -1) {
        return BAD_EXIT;
    }
    
    if (WIFSIGNALED(status)) {
        syslog(LOG_ERR, "supervised process(pid %ld) exited cause of signal %d, status = 0x%X", 
                        (long)pid, WTERMSIG(status), status);
        return RESTART;
    } else if (WIFEXITED(status)) {
        syslog(LOG_ERR, "supervised process(pid %ld) exited with code %d, status = 0x%X", 
                        (long)pid, WTERMSIG(status), status);
        return RESTART;
    }

    /* Ok, can we ever get here ? */

    syslog(LOG_ERR, "waitpid of supervised process(pid %ld) status = 0x%X", 
                    (long)pid, status);
    if (kill(pid, 0) == 0) {
        syslog(LOG_ERR, "supervised process(pid %ld) seems to still be alive, not restarting", 
                        (long)pid);
        goto again;
    } 
    
    return RESTART;
}

static struct sigaction orig_usr1_sa;
static struct sigaction orig_term_sa;
static struct sigaction orig_hup_sa;
static struct sigaction orig_int_sa;

static void restore_signal_handlers(void)
{
    sigaction(SIGUSR1, &orig_usr1_sa, NULL);
    sigaction(SIGTERM, &orig_term_sa, NULL);
    sigaction(SIGHUP,  &orig_hup_sa,  NULL);
    sigaction(SIGINT,  &orig_int_sa,  NULL);
}

static void install_signal_handlers(void)
{
    struct sigaction act;
    memset(&act, 0, sizeof act);

    act.sa_handler = restart_handler;
    sigaction(SIGUSR1, &act, &orig_usr1_sa);

    act.sa_handler = exit_handler;
    sigaction(SIGTERM, &act, &orig_term_sa);
    sigaction(SIGHUP,  &act, &orig_hup_sa);
    sigaction(SIGINT,  &act, &orig_int_sa);
}

//See notes in supervisor.h
int uc_supervise(void)
{
   pid_t pid;

    pid = fork();
    if (pid == -1) {
        syslog(LOG_ERR, "Cannot create supervisor : %s", strerror(errno));
        return -1;
    } else if (pid == 0) {
        return 0;
    }

    //close_files();
    install_signal_handlers();

    syslog(LOG_INFO, "Supervising pid %ld", (long)pid);     

    for (;;) {
        int cmd = watch(pid);
        switch (cmd) {

            case CLEAN_EXIT:
                syslog(LOG_INFO, "Terminating supervised pid %lu", 
                                 (unsigned long)pid);     
                kill_supervised_process(pid);
                wait(NULL);
                syslog(LOG_INFO, "Exiting");
                exit(0);
                break;

            case BAD_EXIT:
                syslog(LOG_ERR, "watch() error, exiting");
                exit(55);
                break;

            case RESTART: {
                pid_t new_pid = -1;

                while (new_pid == -1) {
                    sleep(RESTART_GRACE);
                    syslog(LOG_INFO, "Restarting supervised pid %lu", (unsigned long)pid);     
                    closelog();
                    new_pid = fork();
                    if (new_pid == -1) {
                        syslog(LOG_INFO, "Cannot restart. %s", strerror(errno));
                    } else if (new_pid == 0) { //return so the original process continues 
                        restore_signal_handlers();
                        return 0;
                    }
                }

                pid = new_pid;
                syslog(LOG_INFO, "Supervising new pid %lu", (unsigned long)new_pid);     
                }

                break;

            default:
                exit(38);
                break;
        }
                
    }

    return 0x76;
}