#include <ctype.h>
#include <math.h>
#include <pthread.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#ifdef __OpenBSD__
#define SIG_MIN SIGUSR1
#define SIG_MAX 31
#else
#define SIG_MIN SIGRTMIN
#define SIG_MAX SIGRTMAX
#endif

#ifdef HAS_X11
#include <X11/Xlib.h>
#endif

#define FALSE 0
#define TRUE 1

#define NO_REFRESH 0
#define NO_SIGNAL 0

#define MILISECONDS(t) (t)
#define SECONDS(t) (t * 1000)
#define MINUTES(t) (SECONDS(t) * 60)
#define HOURS(t) (MINUTES(t) * 60)

struct BarModule {
    const char *name;
    void (*init)(void);
    char *(*poll)(void);
    const unsigned long refresh_time;
    const int signal;
    unsigned long last_refresh;
    char *last_text;
    size_t last_text_length;
    int disabled;
};

#include "config.h"

typedef int error_code;

#define ERROR_NONE 0
#define ERROR_UNKNOWN 1
#define ERROR_ARG 2
#define ERROR_MEMORY 3
#define ERROR_MODULE 4
#define ERROR_CONFIG 5
#define ERROR_X11 6

typedef int arg_action;

#define ARG_ACTION_NONE 0
#define ARG_ACTION_CONTINUE 1
#define ARG_ACTION_STOP 2
#define ARG_ACTION_SKIP 3

static const size_t global_prefix_length = sizeof(PREFIX_TEXT) - 1;
static const size_t global_suffix_length = sizeof(SUFFIX_TEXT) - 1;
static const size_t global_seperator_length = sizeof(SEPERATOR_TEXT) - 1;
static const size_t global_module_count =
    sizeof(modules) / sizeof(struct BarModule);
static size_t global_disabled_module_count = 0;
static int global_interpret_flags = TRUE;
static error_code global_last_error = ERROR_NONE;
static const char *global_exec_name;
static struct timespec global_refresh_time;
static clockid_t global_cpu_clock_id;
static unsigned long global_last_refresh;
static pthread_t global_main_thread;
static sigset_t global_real_time_mask;

#ifdef HAS_X11
static Display *global_x_display;
static int global_x_screen;
static Window global_x_root;
static const char *flag_x_display_name;
#elif DEFAULT_SET_ROOT_TITLE != FALSE
#error "DEFAULT_SET_ROOT_TITLE can't be TRUE unless X11 support is included"
#endif

static int flag_print_version = FALSE;
static int flag_print_help = FALSE;
static int flag_print_info = FALSE;
static int flag_set_root_title = DEFAULT_SET_ROOT_TITLE;
static int flag_write_to_stdout = !DEFAULT_SET_ROOT_TITLE;

/* internal function declarations */
static void qtb_exit(void);
static void qtb_exit_with_status(error_code status);
static void check_config(void);
static const char *check_arg_value(char arg, const char *value,
                                   arg_action current_action);
static void set_disabled_modules(const char *arg);
static arg_action interpret_flag_character(char arg, const char *next,
                                           arg_action current_action);
static arg_action interpret_flag_arg(const char *arg, const char *next);
static arg_action interpret_normal_arg(const char *arg, const char *next);
static void interpret_args(int argc, const char **argv);
static void print_flag_error(char arg);
static void print_arg_error(const char *arg);
static void print_help_and_exit(void);
static void print_version_and_exit(void);
static void print_info_and_exit(void);
static void initialize_modules(void);
static unsigned long gcd(unsigned long n1, unsigned long n2);
static unsigned long find_module_refresh_time(void);
static void set_refresh_time(void);
static void signal_handler(int signum);
static void setup_signals(void);
static unsigned long timespec_to_miliseconds(struct timespec *ts);
static void miliseconds_to_timespec(unsigned long ms, struct timespec *out_ts);
static void do_initial_refresh(void);
static void refresh_module(struct BarModule *module);
static int check_and_refresh_module(struct BarModule *module,
                                    unsigned long current_time);
static void module_poll_loop(void);
static void rewrite_bar(void);
static void timespec_diff(const struct timespec *start,
                          const struct timespec *end, struct timespec *out);
static void init_x(void);
static void set_x_root_title(const char *title);

/* module api functions */
void _qtb_internal_log(const char *fmt, ...);
#define qtb_log(...) _qtb_internal_log(__VA_ARGS__)
void qtb_die(void);
void qtb_signal_modules(int sig);
static char *read_file(int fd);
static int handle_execute_parent(pid_t pid, int child_in, int child_out,
                                 int child_err, const char *in,
                                 const char **out, const char **err);
static void handle_execute_child(const char *name, const char **argv, int in,
                                 int out, int err);
void *qtb_malloc(size_t size);
static int is_overflow(size_t n1, size_t n2);
void *qtb_calloc(size_t n, size_t size);
void *qtb_realloc(void *ptr, size_t size);
char *qtb_strdup(const char *str);
void qtb_free(void *ptr);
#ifdef HAS_X11
Display *qtb_get_x_display(void);
#endif

/* function definitions */
int main(int argc, const char **argv) {
    global_exec_name = argv[0];
    global_main_thread = pthread_self();
    check_config();
    interpret_args(argc, argv);
    if (flag_print_help) {
        print_help_and_exit();
    } else if (flag_print_version) {
        print_version_and_exit();
    } else if (flag_print_info) {
        print_info_and_exit();
    } else if (global_disabled_module_count == global_module_count) {
        qtb_log("all modules disabled! exiting...");
        qtb_exit_with_status(ERROR_ARG);
    } else if (global_last_error != ERROR_NONE) {
        qtb_exit();
    }
    if (clock_getcpuclockid(0, &global_cpu_clock_id) != 0) {
        qtb_log("error: could not get cpu clock id");
        qtb_exit_with_status(ERROR_UNKNOWN);
    }
    init_x();
    set_refresh_time();
    setup_signals();
    initialize_modules();
    do_initial_refresh();
    sigprocmask(SIG_UNBLOCK, &global_real_time_mask, NULL);
    module_poll_loop();
    return 0;
}

static void qtb_exit() {
    qtb_exit_with_status(global_last_error);
}

static void qtb_exit_with_status(error_code status) {
    exit(status);
}

static void check_config() {
    if (global_module_count == 0) {
        qtb_log("error: no modules defined");
        qtb_exit_with_status(ERROR_CONFIG);
    }
#ifdef REFRESH_TIME
    if (REFRESH_TIME <= 0) {
        qtb_log("error: refresh time must be greater than 0, was: '%lu'",
                REFRESH_TIME);
        qtb_exit_with_status(ERROR_CONFIG);
    }
#endif
}

/* if current_action is ARG_ACTION_SKIP, print an error and return NULL,
   otherwise, return value */
static const char *check_arg_value(char arg, const char *value,
                                   arg_action current_action) {
    if (current_action == ARG_ACTION_SKIP) {
        qtb_log("error: value already in use by another option: '%s'", arg);
        return NULL;
    }
    if (!value) {
        qtb_log("error: option requires value: '%c'", arg);
    }
    return value;
}

static void set_disabled_modules(const char *arg) {
    if (arg == NULL) {
        return;
    }
    /* reset module disabled state */
    global_disabled_module_count = 0;
    size_t i;
    for (i = 0; i < global_module_count; ++i) {
        modules[i].disabled = FALSE;
    }
    char *arg_copy = qtb_strdup(arg);
    char *end_ptr;
    const char *current = strtok_r(arg_copy, ";", &end_ptr);
    while (current) {
        /* size_t i; */
        for (i = 0; i < global_module_count; ++i) {
            if (strcmp(modules[i].name, current) == 0) {
                modules[i].disabled = TRUE;
                ++global_disabled_module_count;
                goto found_module;
            }
        }
        qtb_log("attempt to disable unknown module: '%s'", current);
    found_module:
        current = strtok_r(NULL, ";", &end_ptr);
        continue;
    }
    qtb_free(arg_copy);
}

/* if current_action is ARG_ACTION_SKIP, error if we encounter an arg with a
 * value */
static arg_action interpret_flag_character(char arg, const char *next,
                                           arg_action current_action) {
    switch (arg) {
    case 'v':
        flag_print_version = TRUE;
        return ARG_ACTION_STOP;
    case 'h':
        flag_print_help = TRUE;
        return ARG_ACTION_STOP;
    case 'i':
        flag_print_info = TRUE;
        return ARG_ACTION_STOP;
    case 's':
        flag_write_to_stdout = TRUE;
        flag_set_root_title = FALSE;
        return ARG_ACTION_CONTINUE;
    case 'n':
        set_disabled_modules(check_arg_value(arg, next, current_action));
        return ARG_ACTION_SKIP;
#ifdef HAS_X11
    case 'd':
        flag_x_display_name = check_arg_value(arg, next, current_action);
        return ARG_ACTION_SKIP;
    case 'r':
        flag_set_root_title = TRUE;
        flag_write_to_stdout = FALSE;
        return ARG_ACTION_CONTINUE;
#else
    case 'd':
    case 'r':
        qtb_log("error: X11 support not included in build");
        global_last_error = ERROR_X11;
        return ARG_ACTION_STOP;
#endif
    default:
        global_last_error = ERROR_ARG;
        print_flag_error(arg);
        return ARG_ACTION_CONTINUE;
    }
}

static arg_action interpret_flag_arg(const char *arg, const char *next) {
    if (!global_interpret_flags) {
        return ARG_ACTION_NONE;
    }
    if (strcmp(arg, "--") == 0) {
        global_interpret_flags = FALSE;
        return ARG_ACTION_NONE;
    }
    if (arg[0] == '\0') {
        print_arg_error(arg);
        return ARG_ACTION_CONTINUE;
    }
    if (arg[0] == '-') {
        if (arg[1] == '\0') {
            print_arg_error(arg);
            return ARG_ACTION_CONTINUE;
        }
        arg_action return_action = ARG_ACTION_CONTINUE;
        ++arg;
        while (*arg) {
            arg_action current_action =
                interpret_flag_character(*arg, next, return_action);
            if (current_action == ARG_ACTION_STOP) {
                return ARG_ACTION_STOP;
            } else if (current_action == ARG_ACTION_CONTINUE && return_action != ARG_ACTION_SKIP) {
                return_action = ARG_ACTION_CONTINUE;
            } else if (current_action == ARG_ACTION_SKIP) {
                return_action = ARG_ACTION_SKIP;
            }
            ++arg;
        }
        return return_action;
    }
    return ARG_ACTION_NONE;
}

static arg_action interpret_normal_arg(const char *arg, const char *next) {
    print_arg_error(arg);
    return ARG_ACTION_CONTINUE;
}

static void interpret_args(int argc, const char **argv) {
    int i;
    for (i = 1; i < argc; ++i) {
        const char *current_arg = argv[i];
        const char *current_value = argv[i + 1];
        arg_action flag_action = interpret_flag_arg(current_arg, current_value);
        switch (flag_action) {
        case ARG_ACTION_NONE:
            break;
        case ARG_ACTION_CONTINUE:
            continue;
        case ARG_ACTION_STOP:
            return;
        case ARG_ACTION_SKIP:
            ++i;
            continue;
        }
        arg_action normal_action =
            interpret_normal_arg(current_arg, current_value);
        switch (normal_action) {
        case ARG_ACTION_NONE:
            break;
        case ARG_ACTION_CONTINUE:
            continue;
        case ARG_ACTION_STOP:
            return;
        case ARG_ACTION_SKIP:
            ++i;
            continue;
        }
    }
}

static void print_flag_error(char arg) {
    if (isprint(arg)) {
        qtb_log("error: unknown option: '%c'", arg);
    } else {
        qtb_log("error: unknown option: '0x%X'", arg);
    }
}

static void print_arg_error(const char *arg) {
    qtb_log("error: unknown argument: '%s'", arg);
}

static void print_help_and_exit() {
    printf("quick-text-bar - a fast and customizable text based bar\n");
#if HAS_X11
    printf("usage: %s [-h] [-v] [-i] [-n MODULES] [-d DISPLAY] [-r] [-s] \n",
           global_exec_name);
#else
    printf("usage: %s [-h] [-v] [-i] [-n MODULES] [-s] \n", global_exec_name);
#endif
    printf("  -h   show this message, then exit\n");
    printf("  -v   print version, then exit\n");
    printf("  -i   print version and module informaton, then exit\n");
    printf("  -n   semi-colon seperated list of modules to disable\n");
#ifdef HAS_X11
    printf("  -d   the X display to use when setting the root window name\n");
    if (DEFAULT_SET_ROOT_TITLE) {
        printf("  -r   set the X root window name when updating the bar "
               "(default)\n");
        printf("  -s   write to stdout when updating the bar\n");
    } else {
        printf("  -r   set the X root window name when updating the bar\n");
        printf("  -s   write to stdout when updating the bar (default)\n");
    }
#else
    printf("  -d   the X display to use when setting the root window name "
           "(disabled)\n");
    printf("  -r   set the X root window name when updating the bar "
           "(disabled)\n");
    printf("  -s   write to stdout when updating the bar\n");
#endif
    qtb_exit();
}

static void print_version_and_exit() {
    printf("quick-text-bar v%s\n", QTB_VERSION);
    qtb_exit();
}

static void print_info_and_exit(void) {
    printf("quick-text-bar v%s\n", QTB_VERSION);
#ifdef HAS_X11
    printf("X11 Support: Present\n");
#else
    printf("X11 Support: Absent\n");
#endif
    printf("Modules: %zu\n", global_module_count);
    size_t i;
    for (i = 0; i < global_module_count; ++i) {
        printf("  %s\n", modules[i].name);
    }
    qtb_exit();
}

static void initialize_modules() {
    int has_error = FALSE;
    size_t i;
    for (i = 0; i < global_module_count; ++i) {
        if (modules[i].disabled) {
            continue;
        }
        if (modules[i].signal + SIG_MIN > SIG_MAX) {
            has_error = TRUE;
            qtb_log("%s: error: signal out of range: %d\n",
                    modules[i].signal + SIG_MIN);
            continue;
        }
        if (modules[i].init) {
            modules[i].init();
        }
        modules[i].last_text = NULL;
        modules[i].last_refresh = 0;
    }
    if (has_error) {
        qtb_log("error: some modules were not initialized, exiting...");
        qtb_exit_with_status(ERROR_MODULE);
    }
}

unsigned long gcd(unsigned long n1, unsigned long n2) {
    if (n1 == 0) {
        return n2;
    }
    return gcd(n2 % n1, n1);
}

unsigned long find_module_refresh_time() {
    size_t i;
    size_t ignore_index;
    size_t result;
    for (i = 0; i < global_module_count; ++i) {
        if (!modules[i].disabled && modules[i].refresh_time != NO_REFRESH) {
            result = modules[i].refresh_time;
            ignore_index = i;
        }
    }
    for (i = 0; i < global_module_count; ++i) {
        if (i == ignore_index) {
            continue;
        }
        if (!modules[i].disabled && modules[i].refresh_time != NO_REFRESH) {
            result = gcd(result, modules[i].refresh_time);
        }
        if (result == 1) {
            break;
        }
    }
    return result;
}

static void set_refresh_time() {
#ifdef REFRESH_TIME
    unsigned long refresh_ms = REFRESH_TIME;
#else
    unsigned long refresh_ms = find_module_refresh_time();
    if (refresh_ms == 0) {     /* if all modules are triggered by signals */
        refresh_ms = HOURS(1); /* refresh once an hour */
    }
#endif
    miliseconds_to_timespec(refresh_ms, &global_refresh_time);
}

static void signal_handler(int signum) {
    struct timespec current_time;
    clock_gettime(global_cpu_clock_id, &current_time);
    int module_updated = FALSE;
    unsigned long current_time_ms = timespec_to_miliseconds(&current_time);
    int adjusted_signum = signum - SIG_MIN;
    size_t i;
    for (i = 0; i < global_module_count; ++i) {
        if (!modules[i].disabled && modules[i].signal == adjusted_signum) {
            refresh_module(&modules[i]);
            modules[i].last_refresh = current_time_ms;
            module_updated = TRUE;
        }
    }
    if (module_updated) {
        rewrite_bar();
    }
}

static void setup_signals() {
    sigemptyset(&global_real_time_mask);
    int i;
    for (i = SIG_MIN; i <= SIG_MAX; ++i) {
        sigaddset(&global_real_time_mask, i);
        struct sigaction action;
        sigemptyset(&action.sa_mask);
        action.sa_flags = 0;
        action.sa_handler = &signal_handler;
        if (sigaction(i, &action, NULL) != 0) {
            qtb_log("error: could not setup handler for signal: %d", i);
            qtb_exit_with_status(ERROR_UNKNOWN);
        }
    }
    sigprocmask(SIG_BLOCK, &global_real_time_mask, NULL);
}

static unsigned long timespec_to_miliseconds(struct timespec *ts) {
    return (ts->tv_sec * 1000) + (ts->tv_nsec / 1000000);
}

static void miliseconds_to_timespec(unsigned long ms, struct timespec *out_ts) {
    out_ts->tv_sec = floorl(ms / 1000);
    out_ts->tv_nsec = (ms % 1000) * 1000000;
}

static void do_initial_refresh() {
    int has_normal_module = FALSE;
    size_t i;
    for (i = 0; i < global_module_count; ++i) {
        if (modules[i].disabled) {
            continue;
        }
        if (modules[i].refresh_time != NO_REFRESH) {
            has_normal_module = TRUE;
        }
        if (modules[i].poll) {
            refresh_module(&modules[i]);
        }
    }
    if (!has_normal_module) {
        rewrite_bar();
    }
}

static void refresh_module(struct BarModule *module) {
    sigprocmask(SIG_BLOCK, &global_real_time_mask, NULL);
    if (module->last_text) {
        qtb_free(module->last_text);
        module->last_text = NULL;
    }
    module->last_text = module->poll();
    if (module->last_text) {
        module->last_text_length = strlen(module->last_text);
    } else {
        module->last_text_length = 0;
    }
    sigprocmask(SIG_UNBLOCK, &global_real_time_mask, NULL);
}

/* return TRUE if the module was refreshed, FALSE otherwise */
static int check_and_refresh_module(struct BarModule *module,
                                    unsigned long current_time) {
    if (module->disabled || module->refresh_time == NO_REFRESH) {
        return FALSE;
    }
    unsigned long module_last_refresh =
        module->last_refresh + module->refresh_time;
    if (module_last_refresh >= global_last_refresh && module->poll) {
        refresh_module(module);
        module->last_refresh = current_time;
        return TRUE;
    }
    return FALSE;
}

static void module_poll_loop() {
    while (TRUE) {
        struct timespec current_time;
        clock_gettime(global_cpu_clock_id, &current_time);
        unsigned long current_time_ms = timespec_to_miliseconds(&current_time);
        int module_refreshed = FALSE;
        size_t i;
        for (i = 0; i < global_module_count; ++i) {
            if (check_and_refresh_module(&modules[i], current_time_ms)) {
                module_refreshed = TRUE;
            }
        }
        if (module_refreshed) {
            rewrite_bar();
        }
        /* calculate how long it took us to run this iteration
           and subtract it from the refresh time */
        struct timespec end_time;
        clock_gettime(global_cpu_clock_id, &end_time);
        struct timespec span_time;
        timespec_diff(&current_time, &end_time, &span_time);
        struct timespec final_sleep_time;
        timespec_diff(&span_time, &global_refresh_time, &final_sleep_time);
        nanosleep(&final_sleep_time, NULL);
    }
}

static void rewrite_bar() {
    /* start with size of prefix, suffix, and null byte */
    size_t bar_length = global_prefix_length + global_suffix_length + 1;
    char *bar_text = qtb_malloc(bar_length);
    bar_text[0] = '\0';
    strcat(bar_text, PREFIX_TEXT);
    size_t i;
    for (i = 0; i < global_module_count - 1; ++i) {
        if (modules[i].last_text) {
            bar_text = qtb_realloc(
                bar_text, (bar_length += modules[i].last_text_length +
                                         global_seperator_length));
            strcat(bar_text, modules[i].last_text);
            strcat(bar_text, SEPERATOR_TEXT);
        }
    }
    struct BarModule *last_module = &modules[global_module_count - 1];
    if (last_module->last_text) {
        bar_text = qtb_realloc(
            bar_text, (bar_length += last_module->last_text_length));
        strcat(bar_text, last_module->last_text);
    }
    strcat(bar_text, SUFFIX_TEXT);
    if (flag_set_root_title) {
        set_x_root_title(bar_text);
    } else {
        printf("%s\n", bar_text);
        fflush(stdout);
    }
    qtb_free(bar_text);
}

/* subtract start from end */
static void timespec_diff(const struct timespec *start,
                          const struct timespec *end, struct timespec *out) {
    out->tv_nsec = end->tv_nsec - start->tv_nsec;
    out->tv_sec = end->tv_sec - start->tv_sec;
    if (out->tv_sec < 0) { /* start > end */
        out->tv_sec = 0;
        out->tv_nsec = 0;
    } else if (out->tv_nsec < 0) {
        out->tv_nsec += 1000000000l;
        --out->tv_sec;
        if (out->tv_sec < 0) {
            out->tv_sec = 0;
        }
    }
}

static void init_x() {
#ifdef HAS_X11
    global_x_display = XOpenDisplay(flag_x_display_name);
    if (!global_x_display) {
        const char *display_name;
        if (flag_x_display_name) {
            display_name = flag_x_display_name;
        } else {
            display_name = XDisplayName(NULL);
        }
        qtb_log("error: failed to open X display: '%s'", display_name);
        if (flag_set_root_title) {
            qtb_exit_with_status(ERROR_X11);
        }
    } else {
        global_x_screen = XDefaultScreen(global_x_display);
        global_x_root = XRootWindow(global_x_display, global_x_screen);
    }
#endif
}

static void set_x_root_title(const char *title) {
#ifdef HAS_X11
    if (global_x_display) {
        XStoreName(global_x_display, global_x_root, title);
        XFlush(global_x_display);
    }
#endif
}

/* module api definitions */
void _qtb_internal_log(const char *fmt, ...) {
    fprintf(stderr, "%s: ", global_exec_name);
    va_list args;
    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
    va_end(args);
    fwrite("\n", 1, 1, stderr);
}

void qtb_die() {
    qtb_exit_with_status(ERROR_MODULE);
}

void qtb_signal_modules(int sig) {
    pthread_kill(global_main_thread, SIG_MIN + sig);
}

void *qtb_malloc(size_t size) {
    return qtb_realloc(NULL, size);
}

static int is_overflow(size_t n1, size_t n2) {
    if (n1 <= 1 || n2 <= 1) {
        return FALSE;
    }
    return (n1 * n2) / n1 != n2;
}

void *qtb_calloc(size_t n, size_t size) {
    if (is_overflow(n, size)) {
        qtb_log(
            "attempt to allocate %zu bytes %zu times failed due to overflow", n,
            size);
        qtb_exit_with_status(ERROR_MEMORY);
    }
    size_t total_size = n * size;
    void *ptr = qtb_malloc(total_size);
    memset(ptr, 0, total_size);
    return ptr;
}

void *qtb_realloc(void *ptr, size_t size) {
    void *new_ptr = realloc(ptr, size);
    if (size != 0 && !new_ptr) {
        qtb_log("error: out of memory");
        qtb_exit_with_status(ERROR_MEMORY);
    }
    return new_ptr;
}

char *qtb_strdup(const char *str) {
    size_t len = strlen(str);
    char *new_str = qtb_malloc(len + 1);
    memcpy(new_str, str, len + 1);
    return new_str;
}

void qtb_free(void *ptr) {
    if (ptr) {
        free(ptr);
    }
}

#ifdef HAS_X11
Display *qtb_get_x_display(void) {
    return global_x_display;
}
#endif