#include "err.h"
#include "library.h"
#include "resource.h"
#include "utils/log.h"

#include <fcntl.h>
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

char *config[CFG_IS_VALID + 1];

void print_help(char *self) {
    LOG_WARN("Usage:");
    LOG_WARN("  %s [options]", self);
    LOG_WARN("Options:");
    LOG_WARN("  --memory_limit     memory limit in MB");
    LOG_WARN("  --nproc_limit      number of processes limit");
    LOG_WARN("  --time_limit       time limit in ms");
    LOG_WARN("  --sandbox_template sandbox template");
    LOG_WARN("  --sandbox_action   sandbox action");
    LOG_WARN("  --uid              user id");
    LOG_WARN("  --gid              group id");
    LOG_WARN("  --file_input       path to input file");
    LOG_WARN("  --file_output      path to output file");
    LOG_WARN("  --file_info        path to info file");
    LOG_WARN("  --program          program to run");
    LOG_WARN("  --program_arg      program argument (only accept one argument)");
    LOG_WARN("  --help             print this help message");
}

void parse(int argc, char *argv[]) {
    static struct option options[] = {
        [CFG_MEMORY_LIMIT]     = {"memory_limit",     required_argument, NULL, 0},
        [CFG_NPROC_LIMIT]      = {"nproc_limit",      required_argument, NULL, 0},
        [CFG_TIME_LIMIT]       = {"time_limit",       required_argument, NULL, 0},
        [CFG_SANDBOX_TEMPLATE] = {"sandbox_template", required_argument, NULL, 0},
        [CFG_SANDBOX_ACTION]   = {"sandbox_action",   required_argument, NULL, 0},
        [CFG_UID]              = {"uid",              optional_argument, NULL, 0},
        [CFG_GID]              = {"gid",              optional_argument, NULL, 0},
        [CFG_FILE_INPUT]       = {"file_input",       required_argument, NULL, 0},
        [CFG_FILE_OUTPUT]      = {"file_output",      required_argument, NULL, 0},
        [CFG_FILE_INFO]        = {"file_info",        required_argument, NULL, 0},
        [CFG_PROGRAM]          = {"program",          required_argument, NULL, 0},
        [CFG_PROGRAM_ARG]      = {"program_arg",      optional_argument, NULL, 0},
        [CFG_IS_VALID]         = {"help",             no_argument,       NULL, 0},
        [CFG_IS_VALID + 1]     = {NULL,               0,                 NULL, 0}
    };

    int c, idx = 0;
    while ((c = getopt_long_only(argc, argv, "", options, &idx)) != -1) {
        if (c != 0) break;

        if (idx < CFG_IS_VALID)
            config[idx] = optarg;
        else if (idx == CFG_IS_VALID) {
            print_help(argv[0]);
            exit(0);
        }
    }

    for (int i = 0; i < CFG_IS_VALID; i++) {
        if (!config[i] && options[i].has_arg == required_argument) {
            print_help(argv[0]);
            LOG_ERR("Missing arguments");
            exit(ERR_ARGUMENTS);
        }
    }

    config[CFG_IS_VALID] = (char *)1;
}

void launch_child() {
    LOG_DEBUG("Launching child process: (%s) (%s)", config[CFG_PROGRAM], config[CFG_PROGRAM_ARG]);

    char *args[] = {config[CFG_PROGRAM], config[CFG_PROGRAM_ARG], NULL};

    /* build stdin */ {
        int fd = open(config[CFG_FILE_INPUT], O_RDONLY);
        dup2(fd, STDIN_FILENO);
        close(fd);
    }

    /* build stdout */ {
        int fd = open(config[CFG_FILE_OUTPUT], O_WRONLY | O_CREAT, 0644);
        dup2(fd, STDOUT_FILENO);
        close(fd);
    }

    setup_all();

    execve(config[CFG_PROGRAM], args, NULL);
}

void dump_info(FILE *dest, struct rusage *usage, int status, long long time_usage) {
    fprintf(dest, "{\"real_time\":%lld,\"cpu_time\":%ld,\"memory\":%ld,", time_usage,
            usage->ru_utime.tv_sec * 1000 + usage->ru_utime.tv_usec / 1000, usage->ru_maxrss);
    if (WIFEXITED(status))
        fprintf(dest, "\"status\":\"exited\",\"code\":%d}", WEXITSTATUS(status));
    else if (WIFSIGNALED(status))
        fprintf(dest, "\"status\":\"killed\",\"code\":%d}", WTERMSIG(status));
    else
        fprintf(dest, "\"status\":\"unknown\",\"code\":0}");
}

int main(int argc, char *argv[]) {
    parse(argc, argv);

    pid_t child = fork();
    if (child < 0) {
        LOG_ERR("Failed to fork (child)");
        exit(ERR_FORK);

    } else if (child == 0) { // Program
        launch_child();
        LOG_ERR("Failed to execute child program");
        exit(ERR_EXEC);

    } else { // Supervisor
        pid_t killer = fork();
        if (killer < 0) {
            LOG_ERR("Failed to fork (killer)");
            exit(ERR_FORK);

        } else if (killer == 0) { // Killer
            long limit = (strtol(config[CFG_TIME_LIMIT], NULL, 10) + 999) / 1000 + 1;
            LOG_INFO("Killer started, time limit: %lds", limit);
            sleep(limit); // 1 more seconds
            LOG_WARN("Killer killed child");
            kill(child, SIGKILL);
            exit(0);

        } else { // stat
            int             status;
            struct rusage   usage;
            pid_t           ret;
            struct timespec time_begin, time_end;

            clock_gettime(CLOCK_MONOTONIC, &time_begin);
            ret = wait4(child, &status, 0, &usage);
            clock_gettime(CLOCK_MONOTONIC, &time_end);

            // LOG_DEBUG("Children exited, status = %d, ret = %d", status, ret);

            long long time_usage =
                (time_end.tv_sec - time_begin.tv_sec) * 1000 + (time_end.tv_nsec - time_begin.tv_nsec) / 1000000;

            kill(killer, SIGKILL);

            if (ret < 0) {
                LOG_ERR("Failed to wait for child process");
                exit(ERR_WAIT);
            }

            FILE *info = fopen(config[CFG_FILE_INFO], "w");
            dump_info(info, &usage, status, time_usage);
            exit(0);
        }
    }
}