#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <string.h>
#include <time.h>
#include <signal.h>
#include <errno.h>
#include <sys/poll.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/un.h>

#include "config.h"

extern float loadAvg();

#define UNIX_PATH_MAX 108

pid_t           child;
pid_t           attach = -1;
pid_t           children[100];
int             copies = 1;
#define         RUNNING 1
#define         STOPPED 2
int             cmdsock = -1;
char            sockname[UNIX_PATH_MAX];

char*
current_time()
{
    time_t      now;
    char        *outline, *p;

    now = time(0);
    outline = ctime(&now);
    if (p = strchr(outline, '\n'))
        *p = 0;

    return outline;
}

void
cleanup()
{
    if (attach == -1)
    {
        if (child)
        {
            /* if i don't send child SIGCONT here, rc5des restarts instead
               of dying.  very weird. */
            kill(-getpgid(child), SIGCONT);
            sleep(1);
            kill(-getpgid(child), SIGTERM);
            sleep(2);
        }
    }
    else
    {
        /* make sure the thing starts running again. */
        kill(-getpgid(attach), SIGCONT);
    }

    if (cmdsock != -1)
    {
        close(cmdsock);
        unlink(sockname);
    }

    exit(0);
}

char *
read_cmd()
{
    int         sock;
    char        *cmd;
    int         size = 1024;
    int         tbr, br;

    if ((sock = accept(cmdsock, 0, 0)) < 0)
    {
        perror("accept");
        return 0;
    }

    tbr = 0;
    cmd = malloc(size);
    if (cmd)
    {
        while (1)
        {
            if ((br = read(sock, cmd + tbr, size - tbr)) < 0)
            {
                switch (errno)
                {
                case EINTR:
                case EAGAIN:
                    continue;

                default:
                    perror("reading command");
                    tbr = 0;
                    break;
                }
            }
            else if (br == 0)
            {
                break;
            }

            tbr += br;

            if (tbr == size)
            {
                char *tmp;
                tmp = realloc(cmd, size *= 2);
                if (!tmp)
                {
                    perror("cannot realloc cmd");
                    tbr = 0;
                    break;
                }
                cmd = tmp;
            }
        }
    }
    else
    {
        cmd = 0;
    }

    close(sock);

    return cmd;
}

void
shepherd(char **argv, int copies)
{
    int i;

    if (setpgid(0,0) == -1)
    {
        fprintf(stderr, "%s (%s): cannot set process group (%s)\n",
                current_time(), getpid(), strerror(errno));
        exit(1);
    }

    for (i = 0; i < copies; i++)
    {
        if ((children[i] = fork()) == 0)
        {
            if (execv(*argv, argv) < 0)
            {
                fprintf(stderr, "%s (%s): error exec'ing command (%s)\n",
                        current_time(), getpid(), strerror(errno));
                exit(1);
            }
        }
    }

    while (1)
    {
        pid_t       status;
        time_t      now;
        char        *p;

        switch(status = waitpid(0, 0, 0))
        {
        case -1:
        case 0:
            break;

        default:
            fprintf(stderr, "%s (%d): child %d finished.\n",
                    current_time(), getpid(), status);
            if (--copies == 0)
                exit(0);
        }
    }
}

void
open_cmd_sock()
{
    if (*sockname)
    {
        struct sockaddr_un addr;

        if ((cmdsock = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
        {
            fprintf(stderr, "%s: cannot create socket (%s)\n",
                    current_time(), strerror(errno));
        }

        addr.sun_family = AF_UNIX;
        strcpy(addr.sun_path, sockname);

        if (cmdsock != -1 &&
            bind(cmdsock, (struct sockaddr *)&addr, sizeof(addr)) < 0)
        {
            fprintf(stderr, "%s: cannot bind socket (%s)\n",
                    current_time(), strerror(errno));
            close(cmdsock);
            cmdsock = -1;
        }

        if (cmdsock != -1 && listen(cmdsock, 5))
        {
            fprintf(stderr, "%s: cannot listen on socket (%s)\n",
                    current_time(), strerror(errno));
            close(cmdsock);
            unlink(sockname);
            cmdsock = -1;
        }
    }
}

void
usage()
{
    fprintf(stderr, "loadwatch [-d <time>] [-h <load>] [-l <load>] [-n <copies] [-p <pid>] [-- <command>]\n");
    fprintf(stderr, "\t-d <int>\tload sampling interval (10 seconds)\n");
    fprintf(stderr, "\t-h <float>\thigh load mark (1.25)\n");
    fprintf(stderr, "\t-l <float>\tlow load mark (0.25)\n");
    fprintf(stderr, "\t-n <copies>\tnumber of children to fork (1)\n");
    fprintf(stderr, "\t-u <filename>\tfile that will be used to externally control a\n\t\t\tloadwatch process.\n");
    fprintf(stderr, "\t-c <command>\tmust be used in conjunction with -u.  use -c to send a\n"
            "\t\t\tcommand to an already running loadwatch.\n"
            "\t\tRUN -> put loadwatch into RUN mode, that is the child process\n"
            "\t\t\truns without regard to load.\n"
            "\t\tSTOP -> put loadwatch into STOP mode, that is, the child\n"
            "\t\t\tprocess will not run.\n"
            "\t\tWATCH -> WATCH mode, the normal loadwatch mode.\n");
    fprintf(stderr, "\t-p <pid>\tpid of process to control (loadwatch will actually\n\t\t\t  send signals to the group containing this pid)\n");
    fprintf(stderr, "\tNOTE: -p and <command> are mutually exclusive, but one has to be\n\t\tspecified.\n");
}

main(int argc, char **argv)
{
    int         delay = 10;
    float       highmark = 1.25;
    float       lowmark = 0.25;
    int         c;
    pid_t       childgroup;

    struct sigaction   handler;
    sockname[0] = 0;

    while ((c = getopt(argc, argv, "d:l:h:p:u:")) != -1)
    {
        switch (c)
        {
        case 'n':
            copies = atoi(optarg);
            break;

        case 'u':
            if (strlen(optarg) >= UNIX_PATH_MAX)
            {
                fprintf(stderr, "maximum path is %d chars long.\n",
                        UNIX_PATH_MAX);
                exit(1);
            }
            strcpy(sockname, optarg);
            break;

        case 'd':
            delay = atoi(optarg);
            break;

        case 'h':
            highmark = atof(optarg);
            break;

        case 'l':
            lowmark = atof(optarg);
            break;

        case 'p':
            attach = getpgid(atoi(optarg));
            break;

        default:
            usage();
            exit(1);
        }
    }

    memset(&handler, 0, sizeof(handler));
    handler.sa_handler = &cleanup;
    sigaction(SIGTERM, &handler, 0);
    sigaction(SIGINT, &handler, 0);
    sigaction(SIGQUIT, &handler, 0);

    if (copies < 1)
    {
        fprintf(stderr, "must run 1 or more copies\n");
        exit(1);
    }

    if (attach == -1 && optind == argc)
    {
        usage();
        printf("current load: %g\n", loadAvg());
        exit(1);
    }

    if (attach == -1 && (child = fork()) == 0)
    {
        shepherd(argv + optind, copies);
    }
    else
    {
        pid_t           status;
        float           avg;
        int             state = RUNNING;

        sleep(1);

        childgroup = attach == -1 ? getpgid(child) : attach;

        if (kill(-childgroup, 0) == -1)
        {
            fprintf(stderr, "%s: cannot kill %d (%s)\n",
                    current_time(), -childgroup, strerror(errno));

            cleanup();
        }

        open_cmd_sock();

        while (1)
        {
            if (attach == -1)
            {
                switch (status = waitpid(child, 0, WNOHANG))
                {
                case -1:
                    perror("problems waiting for child");
                    cleanup();
                    break;

                case 0:
                    break;

                default:
                    fprintf(stderr, "%s: no child process, exiting.\n",
                            current_time());
                    cleanup();
                    break;
                }
            }
            else
            {
                if (kill(-childgroup, 0) == -1)
                {
                    fprintf(stderr, "%s: no attached process, exiting %d.\n",
                            current_time(), -childgroup);
                    cleanup();
                }
            }

            avg = loadAvg();

            if (state == RUNNING && avg >= highmark)
            {
                fprintf(stderr, "%s: load too high, stopping.\n",
                        current_time());
                state = STOPPED;
                kill(-childgroup, SIGSTOP);
            }
            else if (state == STOPPED && avg <= lowmark)
            {
                fprintf(stderr, "%s: load low, continuing.\n", current_time());
                state = RUNNING;
                kill(-childgroup, SIGCONT);
            }

            if (cmdsock == -1)
            {
                sleep(delay);
            }
            else
            {
                struct pollfd   pfd;
                int             pollval;

                pfd.fd = cmdsock;
                pfd.events = POLLIN;

                if ((pollval = poll(&pfd, 1, delay * 1000)) > 0)
                {
                    if (pfd.revents & POLLIN)
                    {
                        char *cmd = read_cmd();

                        if (strcmp(cmd, "RUN") == 0)
                        {
                            fprintf(stderr,
                                    "%s: entering sticky running state.\n",
                                    current_time());
                            state = -RUNNING;
                            kill(-childgroup, SIGCONT);
                        }
                        else if (strcmp(cmd, "STOP") == 0)
                        {
                            fprintf(stderr,
                                    "%s: entering sticky stopped state.\n",
                                    current_time());
                            state = -STOPPED;
                            kill(-childgroup, SIGSTOP);
                        }
                        else if (strcmp(cmd, "WATCH") == 0)
                        {
                            fprintf(stderr,
                                    "%s: entering regular watching state\n",
                                    current_time());
                            if (state < 0)
                                state = -state;
                        }

                        free(cmd);
                    }
                }
                else if (pollval == -1 && errno != EINTR)
                {
                    fprintf(stderr, "%s: select (%s)\n", current_time(),
                            strerror(errno));
                    cleanup();
                }
            }
        }
    }
}

#ifdef linux
float
loadAvg()
{
    float       avg = 0.0;
    FILE        *file;
    char        line[81];

    file = fopen("/proc/loadavg", "r");
    if (fgets(line, 80, file))
    {
        avg = atof(line);
    }
    fclose(file);
    return avg;
}
#elif defined(HAVE_KSTAT)
#include <kstat.h>
#include <sys/param.h>

static kstat_ctl_t *kc;

float loadAvg()
{
    kstat_t       *ks;
    kstat_named_t *kn;

    if (!kc)
    {
        if (!(kc = kstat_open()))
        {
            perror("opening kstat");
            return 0.0;
        }
    }

    ks = kstat_lookup(kc, "unix", 0, "system_misc");
    if (kstat_read(kc, ks, 0) == -1)
    {
        perror("kstat_read");
        return 0.0;
    }

    kn = kstat_data_lookup(ks, "avenrun_1min");
    if (kn)
        return (float)kn->value.ul/FSCALE;
    else
    {
        perror("grabbing avenrun_1min");
        return 0.0;
    }
}

#else
float
loadAvg()
{
    float       avg = 0.0;
    FILE        *file;
    char        line[81], *p;

    file = popen("uptime", "r");
    if (fgets(line, 80, file) && (p = strstr(line, "average:")))
    {
        avg = atof(p + 9);
    }
    pclose(file);
    return avg;
}

#endif
