#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L
#include "yasl/yasl.h"

#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define YASL_LOOP_INIT_CAP 32

typedef int (*yasl_loop_fd_fn)(int fd, short revents, void *ud);

struct yasl_fd_source {
    yasl_loop_fd_fn fn;
    void *ud;
    int id;
    int fd;
    short events;
};

struct yasl_loop {
    struct yasl_fd_source *fds;
    struct pollfd *pollfds;
    int count;
    int capacity;
    int next_id;
    bool running;
    int wake_rfd;
    int wake_wfd;
    int wake_watch;
};

uint64_t yasl_loop_now_ns(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ull + (uint64_t)ts.tv_nsec;
}

static int loop_wake_cb(int fd, short revents, void *ud)
{
    (void)revents;
    (void)ud;
    char buf[128];
    ssize_t n;
    while ((n = read(fd, buf, sizeof(buf))) > 0 || (n < 0 && errno == EINTR)) {}
    return 0;
}

struct yasl_loop *yasl_loop_create(void)
{
    struct yasl_loop *l = calloc(1, sizeof(*l));
    if (!l) return NULL;

    l->capacity = YASL_LOOP_INIT_CAP;
    l->fds = malloc((size_t)l->capacity * sizeof(*l->fds));
    l->pollfds = malloc((size_t)l->capacity * sizeof(*l->pollfds));
    if (!l->fds || !l->pollfds) {
        free(l->fds);
        free(l->pollfds);
        free(l);
        return NULL;
    }

    l->next_id = 1;
    l->wake_rfd = -1;
    l->wake_wfd = -1;

    int p[2];
#if defined(__linux__) && defined(O_CLOEXEC)
    if (pipe2(p, O_CLOEXEC | O_NONBLOCK) == 0) {
        l->wake_rfd = p[0];
        l->wake_wfd = p[1];
    }
#else
    if (pipe(p) == 0) {
        int f0 = fcntl(p[0], F_GETFL, 0);
        int f1 = fcntl(p[1], F_GETFL, 0);
        if (f0 >= 0) fcntl(p[0], F_SETFL, f0 | O_NONBLOCK);
        if (f1 >= 0) fcntl(p[1], F_SETFL, f1 | O_NONBLOCK);
        fcntl(p[0], F_SETFD, FD_CLOEXEC);
        fcntl(p[1], F_SETFD, FD_CLOEXEC);
        l->wake_rfd = p[0];
        l->wake_wfd = p[1];
    }
#endif

    if (l->wake_rfd >= 0) {
        l->wake_watch = yasl_loop_add_fd(l, l->wake_rfd, POLLIN, loop_wake_cb, l);
    }
    return l;
}

void yasl_loop_destroy(struct yasl_loop *l)
{
    if (!l) return;
    if (l->wake_watch > 0)
        yasl_loop_del_fd(l, l->wake_watch);
    if (l->wake_rfd >= 0)
        close(l->wake_rfd);
    if (l->wake_wfd >= 0)
        close(l->wake_wfd);
    free(l->fds);
    free(l->pollfds);
    free(l);
}

int yasl_loop_add_fd(struct yasl_loop *l, int fd, short events,
                     int (*fn)(int fd, short revents, void *ud), void *ud)
{
    if (!l || fd < 0)
        return -1;

    if (l->count >= l->capacity) {
        int new_cap = l->capacity ? l->capacity * 2 : YASL_LOOP_INIT_CAP;
        struct yasl_fd_source *new_fds = realloc(l->fds, (size_t)new_cap * sizeof(*new_fds));
        if (!new_fds) return -1;
        l->fds = new_fds;

        struct pollfd *new_pollfds = realloc(l->pollfds, (size_t)new_cap * sizeof(*new_pollfds));
        if (!new_pollfds) return -1;
        l->pollfds = new_pollfds;
        l->capacity = new_cap;
    }

    int id = l->next_id++;
    if (l->next_id <= 0) l->next_id = 1;

    l->fds[l->count] = (struct yasl_fd_source){
        .fn = fn, .ud = ud, .id = id, .fd = fd, .events = events,
    };
    l->pollfds[l->count] = (struct pollfd){
        .fd = fd, .events = events, .revents = 0,
    };
    l->count++;
    return id;
}

int yasl_loop_del_fd(struct yasl_loop *l, int id)
{
    if (!l || id <= 0) return -1;
    for (int i = 0; i < l->count; i++) {
        if (l->fds[i].id == id) {
            l->count--;
            l->fds[i] = l->fds[l->count];
            l->pollfds[i] = l->pollfds[l->count];
            return 0;
        }
    }
    return -1;
}

int yasl_loop_set_fd_events(struct yasl_loop *l, int id, short events)
{
    if (!l || id <= 0) return -1;
    for (int i = 0; i < l->count; i++) {
        if (l->fds[i].id == id) {
            l->fds[i].events = events;
            l->pollfds[i].events = events;
            return 0;
        }
    }
    return -1;
}

void yasl_loop_kick(struct yasl_loop *l)
{
    if (!l || l->wake_wfd < 0) return;
    char b = 1;
    ssize_t n;
    do {
        n = write(l->wake_wfd, &b, 1);
    } while (n < 0 && errno == EINTR);
}

int yasl_loop_step(struct yasl_loop *l, int timeout_ms)
{
    if (!l) return -1;
    if (l->count <= 0) {
        if (timeout_ms > 0) {
            struct timespec req = {
                .tv_sec = timeout_ms / 1000,
                .tv_nsec = (long)(timeout_ms % 1000) * 1000000L,
            };
            nanosleep(&req, NULL);
        }
        return 0;
    }

    int n;
    do {
        n = poll(l->pollfds, (nfds_t)l->count, timeout_ms);
    } while (n < 0 && errno == EINTR);

    if (n < 0) return -1;
    if (n == 0) return 0;

    int handled = 0;
    int count = l->count;

    struct yasl_fd_source sstack[32];
    struct pollfd pstack[32];
    struct yasl_fd_source *scopy = (count <= 32) ? sstack : malloc((size_t)count * sizeof(*scopy));
    struct pollfd *pcopy = (count <= 32) ? pstack : malloc((size_t)count * sizeof(*pcopy));
    if (!scopy || !pcopy) {
        if (scopy && scopy != sstack) free(scopy);
        if (pcopy && pcopy != pstack) free(pcopy);
        return -1;
    }

    memcpy(scopy, l->fds, (size_t)count * sizeof(*scopy));
    memcpy(pcopy, l->pollfds, (size_t)count * sizeof(*pcopy));

    for (int i = 0; i < count; i++) {
        if (pcopy[i].revents) {
            short revents = pcopy[i].revents;
            // Verify that the source was not deleted by an earlier callback in this step
            bool valid = false;
            for (int j = 0; j < l->count; j++) {
                if (l->fds[j].id == scopy[i].id) {
                    l->pollfds[j].revents = 0;
                    valid = true;
                    break;
                }
            }
            if (valid && scopy[i].fn) {
                scopy[i].fn(scopy[i].fd, revents, scopy[i].ud);
                handled++;
            }
        }
    }

    if (scopy != sstack) free(scopy);
    if (pcopy != pstack) free(pcopy);
    return handled;
}

int yasl_loop_run(struct yasl_loop *l)
{
    if (!l) return -1;
    l->running = true;
    while (l->running) {
        if (yasl_loop_step(l, -1) < 0)
            return -1;
    }
    return 0;
}

void yasl_loop_quit(struct yasl_loop *l)
{
    if (l) l->running = false;
}

bool yasl_loop_is_running(const struct yasl_loop *l)
{
    return l ? l->running : false;
}