#define AL_LOG_SECTION "tarod"
#include <al/log.h>
#include <al/array.h>
#include <al/str.h>
#include <nnwt/common.h>
#include <nnwt/event_loop.h>
#include <nnwt/signal.h>
#include <nnwt/fs_event.h>
#include <nnwt/error.h>
#include <nnwt/socket.h>
#include <stl/window.h>
#include <jansson.h>
#include <sys/wait.h>

#include "file_ext.h"
#include "config.h"

//#define USE_SWAYIPC

struct monitor {
    struct stl_monitor *os;
    str config_dir;
    str selection_path;
    struct nn_fs_event selection_event;
    bool self_event_skip;
    pid_t pid;
    bool paused;
    bool covered;
    struct daemon *dmon;
};

struct daemon {
    struct nn_event_loop loop;
    struct nn_signal quit_signal;
    s32 context_poll_fd;
    struct nn_poll context_poll;
    array(struct monitor *) monitors;
#ifdef USE_SWAYIPC
#define SWAY_IPC_HEADER_LENGTH 14
#define SWAY_IPC_SUBSCRIBE 2
    struct nn_socket ipc_socket;
    bool ipc_connected;
    bool ipc_errored;
    struct nn_poll ipc_poll;
    struct {
        u8 *buf;
        ssize_t index;
        ssize_t size;
    } out;
    struct {
        u8 *buf;
        ssize_t index;
        ssize_t size;
        ssize_t alloc;
    } ipc;
#endif
    struct config conf;
    struct nn_fs_event conf_dir_event;
    struct nn_fs_event conf_event;
};

static s32 parse_and_set_props(str *selection, str *wallpaper_id, bool *is_package)
{
    u32 index = al_str_find(selection, '|');
    u32 rindex = al_str_rfind(selection, '|');
    if (index == AL_STR_NO_POS || rindex == AL_STR_NO_POS || index == rindex) {
        return -1;
    }

    // index < 2 or id->length == 0 is expected when a monitor has no wallpaper set.
    if (index < 2) return 1;
    str id = al_str_substr(selection, 0, index - 2);
    str type = al_str_substr(selection, index - 2, index);
    if (id.length == 0) return 1;

    al_str_clone(wallpaper_id, &id);
    if (al_str_eq(&type, &al_str_c("_p"))) {
        *is_package = true;
    } else if (al_str_eq(&type, &al_str_c("_d"))) {
        *is_package = false;
    } else {
        return -1;
    }

    return 0;
}

static bool send_signal_to_pid(pid_t pid, s32 signal)
{
    if (kill(pid, signal) != 0) {
        log_error("kill(%d) failed (%s)", pid, nn_strerror(errno));
        return false;
    }
    return true;
}

static bool signal_monitor(struct monitor *monitor, s32 signal)
{
    if (monitor->pid == -1) return false;
    if (!send_signal_to_pid(monitor->pid, signal)) {
        monitor->pid = -1;
        return false;
    }
    return true;
}

static pid_t exec_or_exit(char *args[])
{
    pid_t pid = fork();
    if (pid == -1) {
        log_error("fatal: failed to fork process, exiting...");
        exit(EXIT_FAILURE);
    } else if (pid == 0) {
        if (execvp(args[0], args) == -1) {
            log_error("fatal: failed to run '%s', exiting...", args[0]);
            // This should return the pid of init if the parent process has already exited.
            pid_t ppid = getppid();
            if (ppid != 1) send_signal_to_pid(ppid, SIGINT);
            // Rely on maybe_kill_wallpaper() to kill this fork.
        }
    }
    return pid;
}

static void maybe_kill_wallpaper(struct monitor *monitor)
{
    if (monitor->pid == -1) return;
    if (monitor->paused) {
        signal_monitor(monitor, SIGCONT);
        monitor->paused = false;
    }
    signal_monitor(monitor, SIGTERM);
    monitor->pid = -1;
}

static bool search_wallpaper_dirs(struct daemon *dmon, str *id, str *path)
{
    bool found = false;
    log_info("searching for wallpaper of id %.*s", al_str_x(id));
    str *dir;
    al_array_foreach_ptr(dmon->conf.wallpaper_dirs, i, dir) {
        log_info("checking dir %.*s", al_str_x(dir));
        al_str_clone(path, dir);
        if (al_str_at(path, 1) != '/') al_str_cat(path, &al_str_c("/"));
        al_str_cat(path, id);
        if (nn_dir_exists(path)) {
            log_info("found wallpaper");
            found = true;
            break;
        }
        al_str_free(path);
    }
    if (!found) log_error("failed to find wallpaper of id %.*s", al_str_x(id));
    al_str_free(id);
    return found;
}

static bool run_wallpaper(struct daemon *dmon, struct monitor *monitor)
{
    maybe_kill_wallpaper(monitor);

    if (dmon->conf.newly_created) {
        al_assert(!dmon->conf.valid);
        return false;
    }

    if (!dmon->conf.valid) {
        log_error("config invalid, not setting wallpaper on %s", monitor->os->name);
        return false;
    }

    str selection;
    monitor->self_event_skip = true;
    if (!nn_file_open_and_read_wholly(&monitor->selection_path, &selection, true)) {
        monitor->self_event_skip = false;
        return false;
    }

    if (selection.length == 0) {
        log_error("config for %s is empty", monitor->os->name);
        al_str_free(&selection);
        return false;
    }

    str wallpaper_id = al_str_null();
    bool is_package;
    s32 ret = parse_and_set_props(&selection, &wallpaper_id, &is_package);
    al_str_free(&selection);
    if (ret != 0) {
        if (ret == -1) {
            log_error("failed to set properties for wallpaper on %s", monitor->os->name);
        }
        al_str_free(&wallpaper_id);
        return false;
    }

    str wallpaper_path = al_str_null();
    if (!search_wallpaper_dirs(dmon, &wallpaper_id, &wallpaper_path)) {
        return false;
    }

    if (is_package) {
        al_str_cat(&wallpaper_path, &al_str_c("/scene.pkg"));
        if (!nn_file_exists(&wallpaper_path)) {
            al_str_insert(&wallpaper_path, wallpaper_path.length - 9, &al_str_c("gif"));
            if (!nn_file_exists(&wallpaper_path)) {
                log_error("wallpaper at %.*s has no package file", al_str_x(&wallpaper_path));
                al_str_free(&wallpaper_path);
                return false;
            }
        }
    }

    log_info("setting wallpaper %.*s on %s", al_str_x(&wallpaper_path), monitor->os->name);

    char width[8], height[8];
    al_snprintf(width, sizeof(width), "%d", monitor->os->width / monitor->os->scale);
    al_snprintf(height, sizeof(height), "%d", monitor->os->height / monitor->os->scale);

    char *c_str0 = al_str_to_c_str(&wallpaper_path);
    char *c_str1 = al_str_to_c_str(&dmon->conf.asset_dir);
    if (is_package) {
        monitor->pid = exec_or_exit((char *[]){
            "mauri", "-p", c_str0, "-a", c_str1, "-m", monitor->os->name, "-w", width, "-h", height, NULL });
    } else {
        monitor->pid = exec_or_exit((char *[]){
            "mauri", "-d", c_str0, "-a", c_str1, "-m", monitor->os->name, "-w", width, "-h", height, NULL });
    }
    al_str_free(&wallpaper_path);
    al_free(c_str0);
    al_free(c_str1);

    if (monitor->covered && signal_monitor(monitor, SIGSTOP)) {
        monitor->paused = true;
    }

    return true;
}

static void selection_fs_event_callback(void *userdata, struct inotify_event *event)
{
    struct monitor *monitor = (struct monitor *)userdata;
    (void)event;
    if (monitor->self_event_skip) {
        monitor->self_event_skip = false;
        return;
    }
    struct daemon *dmon = monitor->dmon;
    run_wallpaper(dmon, monitor);
}

static void create_monitor(struct daemon *dmon, str *dir, struct stl_monitor *os_monitor)
{
    struct monitor *monitor = al_alloc_object(struct monitor);
    monitor->os = os_monitor;
    monitor->dmon = dmon;
    al_array_push(dmon->monitors, monitor);

    al_str_clone(&monitor->config_dir, dir);

    al_str_clone(&monitor->selection_path, &monitor->config_dir);
    al_str_cat(&monitor->selection_path, &al_str_c("/selection"));

    if (!nn_file_exists(&monitor->selection_path)) {
        nn_file_create(&monitor->selection_path);
        nn_file_open_and_replace(&monitor->selection_path, &al_str_c("||"), false);
    }

    monitor->pid = -1;
    monitor->paused = false;
    monitor->covered = false;

    run_wallpaper(dmon, monitor);

    monitor->self_event_skip = false;
    nn_fs_event_init(&monitor->selection_event, &monitor->selection_path,
        NNWT_FS_CLOSE_WRITE, selection_fs_event_callback, monitor);
    nn_fs_event_start(&monitor->selection_event, &monitor->dmon->loop);
}

static void output_discovered_callback(void *userdata, struct stl_monitor *os_monitor)
{
    struct daemon *dmon = (struct daemon *)userdata;

    log_info("Output discovered: %s (%ux%u) (vertical: %s).", os_monitor->name,
        os_monitor->width, os_monitor->height, BOOLSTR(os_monitor->vertical));

    struct monitor *monitor;
    al_array_foreach(dmon->monitors, i, monitor) {
        if (strcmp(monitor->os->name, os_monitor->name) == 0) {
            if (monitor->pid == -1) {
                run_wallpaper(dmon, monitor);
            }
            return;
        }
    }

    str monitor_dir;
    al_str_clone(&monitor_dir, &dmon->conf.config_dir);
    al_str_cat(&monitor_dir, &al_str_c("/"));
    al_str_cat(&monitor_dir, &al_str_cr(os_monitor->name));

    if (nn_dir_exists(&monitor_dir) || nn_dir_create(&monitor_dir)) {
        create_monitor(dmon, &monitor_dir, os_monitor);
    }

    al_str_free(&monitor_dir);
}

// https://vi.stackexchange.com/questions/25030/vim-not-firing-inotify-events-when-writing-file
static void conf_dir_fs_event_callback(void *userdata, struct inotify_event *event)
{
    struct daemon *dmon = (struct daemon *)userdata;
    if (al_strncmp(event->name, "config.json", event->len) == 0) {
        nn_fs_event_stop(&dmon->conf_event);
        nn_fs_event_start(&dmon->conf_event, &dmon->loop);
    }
}

static void conf_fs_event_callback(void *userdata, struct inotify_event *event)
{
    struct daemon *dmon = (struct daemon *)userdata;
    (void)event;
    if (!dmon->conf.errored) close_config(&dmon->conf);
    if (open_config(&dmon->conf, false) == CONFIG_LOADED && dmon->conf.valid) {
        log_info("config successfully reloaded");
    } else {
        log_error("failed to reload config");
        return;
    }
    struct monitor *monitor;
    al_array_foreach(dmon->monitors, i, monitor) {
        run_wallpaper(dmon, monitor);
    }
}

static void close_monitor(struct monitor *monitor)
{
    al_str_free(&monitor->config_dir);
    al_str_free(&monitor->selection_path);
    nn_fs_event_free(&monitor->selection_event);
}

#ifdef USE_SWAYIPC
static void ipc_handle_subscribe_success(json_t *root)
{
    if (json_boolean_value(root)) {
        log_info("successfully subscribed to sway background covered events");
    } else {
        log_info("failed to subscribe to sway background covered events");
    }
}

static void ipc_handle_response(struct daemon *dmon, u8 *buf, u32 size)
{
    log_debug("IN: %.*s", size, (char *)buf);

    json_error_t error;
    json_t *root = json_loadb((char *)buf, size, 0, &error);
    if (!root) {
        log_error("failed to parse sway-ipc json payload (%s)", error.text);
        goto done;
    }

    json_t *success = json_object_get(root, "success");
    if (success) {
        ipc_handle_subscribe_success(success);
        goto done;
    }

    json_t *output = json_object_get(root, "output");
    if (output) {
        bool covered = json_boolean_value(json_object_get(root, "covered"));
        struct monitor *monitor;
        al_array_foreach(dmon->monitors, i, monitor) {
            if (strcmp(monitor->os->name, json_string_value(output)) != 0) continue;
            monitor->covered = covered;
            if (monitor->pid == -1) break;
            if (covered && !monitor->paused && signal_monitor(monitor, SIGSTOP)) {
                monitor->paused = true;
            } else if (!covered && monitor->paused && signal_monitor(monitor, SIGCONT)) {
                monitor->paused = false;
            }
            break;
        }
    }

done:
    if (root) json_decref(root);
}

static bool ipc_prepare_payload(struct daemon *dmon)
{
    json_t *payload = json_array();
    json_array_append_new(payload, json_string_nocheck("background_covered"));
    size_t length = json_dumpb(payload, NULL, 0, 0);
    if (length == 0) return false;
    size_t bufsize = SWAY_IPC_HEADER_LENGTH + length;
    dmon->out.buf = al_malloc(bufsize);
    al_memcpy(dmon->out.buf, "i3-ipc", al_strlen("i3-ipc"));
    al_memcpy(dmon->out.buf + 6, &length, sizeof(s32));
    s32 type = SWAY_IPC_SUBSCRIBE;
    al_memcpy(dmon->out.buf + 10, &type, sizeof(s32));
    length = json_dumpb(payload, (char *)dmon->out.buf + SWAY_IPC_HEADER_LENGTH, bufsize - SWAY_IPC_HEADER_LENGTH, 0);
    json_decref(payload);
    dmon->out.size = bufsize;
    dmon->out.index = 0;
    return true;
}

static void ensure_ipc_buf_size(struct daemon *dmon, ssize_t size)
{
    if (dmon->ipc.alloc < size) {
        dmon->ipc.buf = al_realloc(dmon->ipc.buf, size);
        dmon->ipc.alloc = size;
    }
}

static void ipc_poll_callback(void *userdata, s32 revents)
{
    struct daemon *dmon = (struct daemon *)userdata;

    if (revents & NNWT_POLL_WRITE) {
        al_assert(!dmon->ipc_connected);
        if (!dmon->ipc_connected) {
            log_info("connected to sway ipc socket");
            dmon->ipc_connected = true;
        }
        ssize_t ret = nn_socket_write(&dmon->ipc_socket, dmon->out.buf + dmon->out.index, dmon->out.size - dmon->out.index);
        if (ret <= 0) {
            nn_signal_send(&dmon->quit_signal);
            return;
        }
        dmon->out.index += ret;
        if (dmon->out.index == dmon->out.size) {
            nn_poll_stop(&dmon->ipc_poll);
            nn_poll_set(&dmon->ipc_poll, nn_socket_get_fd(&dmon->ipc_socket), NNWT_POLL_READ);
            nn_poll_start(&dmon->ipc_poll, &dmon->loop);
        }
    }

    if (revents & NNWT_POLL_READ) {
        if (dmon->ipc.index < SWAY_IPC_HEADER_LENGTH) {
            ssize_t ret = nn_socket_read(&dmon->ipc_socket, dmon->ipc.buf + dmon->ipc.index, SWAY_IPC_HEADER_LENGTH - dmon->ipc.index);
            if (ret <= 0) {
                nn_signal_send(&dmon->quit_signal);
                return;
            }
            dmon->ipc.index += ret;
        }
        if (!dmon->ipc.size && dmon->ipc.index == SWAY_IPC_HEADER_LENGTH) {
            dmon->ipc.size = *(s32 *)(dmon->ipc.buf + 6) + SWAY_IPC_HEADER_LENGTH;
            ensure_ipc_buf_size(dmon, dmon->ipc.size);
        }
        if (dmon->ipc.size) {
            ssize_t ret = nn_socket_read(&dmon->ipc_socket, dmon->ipc.buf + dmon->ipc.index, dmon->ipc.size - dmon->ipc.index);
            if (ret <= 0) {
                nn_signal_send(&dmon->quit_signal);
                return;
            }
            dmon->ipc.index += ret;
            if (dmon->ipc.index == dmon->ipc.size) {
                ipc_handle_response(dmon, dmon->ipc.buf + SWAY_IPC_HEADER_LENGTH, dmon->ipc.size - SWAY_IPC_HEADER_LENGTH);
                dmon->ipc.size = 0;
                dmon->ipc.index = 0;
            }
        }
    }
}
#endif

static void context_poll_callback(void *userdata, s32 revents)
{
    (void)userdata;
    al_assert(revents & NNWT_POLL_READ);
    stl_global_process_events();
}

static void quit_signal_callback(void *userdata)
{
    struct daemon *dmon = (struct daemon *)userdata;
    struct monitor *monitor;
    al_array_foreach(dmon->monitors, i, monitor) {
        nn_fs_event_stop(&monitor->selection_event);
    }
    nn_fs_event_stop(&dmon->conf_dir_event);
    nn_fs_event_stop(&dmon->conf_event);
#ifdef USE_SWAYIPC
    if (!dmon->ipc_errored) nn_poll_stop(&dmon->ipc_poll);
#endif
    if (dmon->context_poll_fd != -1) nn_poll_stop(&dmon->context_poll);
    nn_signal_stop(&dmon->quit_signal);
}

static struct daemon dmon = { 0 };

// https://stackoverflow.com/a/59071370
static void sigchld_handler(s32 signum)
{
    (void)signum;
    s32 saved_errno = errno;
    while (waitpid((pid_t)(-1), 0, WNOHANG) > 0) {}
    errno = saved_errno;
}

static void sigint_handler(s32 signum)
{
    (void)signum;
    nn_signal_send(&dmon.quit_signal);
}

s32 main()
{
    s32 exit_status = EXIT_FAILURE;

    if (!nn_common_init(NULL)) return exit_status;

    switch (open_config(&dmon.conf, true)) {
    case CONFIG_CREATED:
        log_info("config generated at $HOME%s/config.json, go fill it out", CONFIG_PATH);
        break;
    case CONFIG_LOADED:
        break;
    case CONFIG_ERRORED:
        log_error("failed to open config");
        nn_common_close();
        return exit_status;
    }

    signal(SIGINT, sigint_handler);
    signal(SIGCHLD, sigchld_handler);

    nn_event_loop_init(&dmon.loop);

    nn_signal_init(&dmon.quit_signal, &dmon.loop, quit_signal_callback, &dmon);
    nn_signal_start(&dmon.quit_signal);

    al_array_init(dmon.monitors);
    stl_set_output_discovered_callback(output_discovered_callback, &dmon);
    if (!stl_global_init(true)) goto out;
    dmon.context_poll_fd = stl_global_get_poll_fd();
    if (dmon.context_poll_fd != -1) {
        nn_poll_init(&dmon.context_poll, context_poll_callback, &dmon);
        nn_poll_set(&dmon.context_poll, dmon.context_poll_fd, NNWT_POLL_READ);
        nn_poll_start(&dmon.context_poll, &dmon.loop);
    }

    str monitors;
    al_str_from(&monitors, "");
    struct monitor *monitor;
    al_array_foreach(dmon.monitors, i, monitor) {
        al_str_cat(&monitors, &al_str_cr(monitor->os->name));
        al_str_cat(&monitors, &al_str_c("\n"));
    }

    str monitors_path;
    al_str_clone(&monitors_path, &dmon.conf.config_dir);
    al_str_cat(&monitors_path, &al_str_c("/monitors"));
    bool monitors_exists = nn_file_exists(&monitors_path) || nn_file_create(&monitors_path);
    if (monitors_exists) {
        nn_file_open_and_replace(&monitors_path, &monitors, false);
    }
    al_str_free(&monitors_path);
    al_str_free(&monitors);
    if (!monitors_exists) goto out;

#ifdef USE_SWAYIPC
    dmon.ipc.buf = NULL;
    dmon.ipc.size = 0;
    dmon.ipc.index = 0;
    dmon.ipc.alloc = 0;
    ensure_ipc_buf_size(&dmon, SWAY_IPC_HEADER_LENGTH);
    ipc_prepare_payload(&dmon);
    dmon.ipc_connected = false;
    dmon.ipc_errored = false;
    dmon.ipc_socket.type = NNWT_SOCKET_UNIX;
    nn_socket_init(&dmon.ipc_socket, NNWT_SOCKET_NONBLOCKING);
    nn_poll_init(&dmon.ipc_poll, ipc_poll_callback, &dmon);
    nn_poll_set(&dmon.ipc_poll, nn_socket_get_fd(&dmon.ipc_socket), NNWT_POLL_WRITE);
    char *swaysock = getenv("SWAYSOCK");
    if (swaysock && nn_socket_connect(&dmon.ipc_socket, &al_str_cr(swaysock), 0)) {
        nn_poll_start(&dmon.ipc_poll, &dmon.loop);
    } else {
        dmon.ipc_errored = true;
        log_warn("couldn't connect to sway ipc socket");
    }
#endif

    nn_fs_event_init(&dmon.conf_dir_event, &dmon.conf.config_dir,
        NNWT_FS_CREATE, conf_dir_fs_event_callback, &dmon);
    nn_fs_event_start(&dmon.conf_dir_event, &dmon.loop);
    nn_fs_event_init(&dmon.conf_event, &dmon.conf.config_path,
        NNWT_FS_CLOSE_WRITE, conf_fs_event_callback, &dmon);
    nn_fs_event_start(&dmon.conf_event, &dmon.loop);

    nn_event_loop_run(&dmon.loop);

    exit_status = EXIT_SUCCESS;

    nn_fs_event_free(&dmon.conf_dir_event);
    nn_fs_event_free(&dmon.conf_event);

#ifdef USE_SWAYIPC
    al_free(dmon.out.buf);
    al_free(dmon.ipc.buf);
    nn_socket_close(&dmon.ipc_socket);
#endif

out:
    al_array_foreach(dmon.monitors, i, monitor) {
        maybe_kill_wallpaper(monitor);
        close_monitor(monitor);
        al_free(monitor);
    }
    al_array_free(dmon.monitors);

    nn_event_loop_destroy(&dmon.loop);

    close_config(&dmon.conf);

    stl_global_close();
    nn_common_close();

    return exit_status;
}