// uasync.c

#include "u_async.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>

#ifndef FD_SETSIZE
#define FD_SETSIZE 1024  // Assume standard size; adjust if needed for your platform
#endif

// Timeout node
struct timeout_node {
    struct timeval expire_time;
    void* arg;
    timeout_callback_t callback;
    struct timeout_node* next;
    struct timeout_node* prev;  // Doubly linked for easier removal
};

// Socket node
struct socket_node {
    int fd;
    socket_callback_t read_cbk;
    socket_callback_t write_cbk;
    socket_callback_t except_cbk;
    void* user_data;
    struct socket_node* next;
};

// Global state
static struct timeout_node* timeout_head = NULL;  // Sorted by expire_time, earliest first
static struct socket_node* socket_head = NULL;
static int max_fd = -1;

// New: Persistent master fd_sets, updated only on add/remove
static fd_set master_readfds;
static fd_set master_writefds;
static fd_set master_exceptfds;

// New: FD-to-node map for faster post-select lookup (addresses point 2)
static struct socket_node* fd_to_node[FD_SETSIZE];

// Helper to get current time
static void get_current_time(struct timeval* tv) {
    gettimeofday(tv, NULL);
}

// Helper to compare times: returns 1 if a > b, 0 if equal, -1 if a < b
static int timeval_cmp(const struct timeval* a, const struct timeval* b) {
    if (a->tv_sec > b->tv_sec) return 1;
    if (a->tv_sec < b->tv_sec) return -1;
    if (a->tv_usec > b->tv_usec) return 1;
    if (a->tv_usec < b->tv_usec) return -1;
    return 0;
}

// Helper to add timeval: tv += dt (timebase units)
static void timeval_add_tb(struct timeval* tv, int dt) {
    tv->tv_usec += (dt % 10000) * 100;
    tv->tv_sec += dt / 10000 + tv->tv_usec / 1000000;
    tv->tv_usec %= 1000000;
}

// Insert timeout into sorted list
static struct timeout_node* insert_timeout(struct timeout_node* new_node) {
    if (!timeout_head || timeval_cmp(&new_node->expire_time, &timeout_head->expire_time) < 0) {
        new_node->next = timeout_head;
        new_node->prev = NULL;
        if (timeout_head) timeout_head->prev = new_node;
        timeout_head = new_node;
        return new_node;
    }

    struct timeout_node* cur = timeout_head;
    while (cur->next && timeval_cmp(&new_node->expire_time, &cur->next->expire_time) >= 0) {
        cur = cur->next;
    }
    new_node->next = cur->next;
    new_node->prev = cur;
    if (cur->next) cur->next->prev = new_node;
    cur->next = new_node;
    return new_node;
}

// Process expired timeouts
static void process_timeouts() {
    struct timeval now;
    get_current_time(&now);

    while (timeout_head && timeval_cmp(&timeout_head->expire_time, &now) <= 0) {
        struct timeout_node* node = timeout_head;
        timeout_head = node->next;
        if (timeout_head) timeout_head->prev = NULL;

        if (node->callback) node->callback(node->arg);
        free(node);
    }
}

// Compute time to next timeout
static void get_next_timeout(struct timeval* tv) {
    struct timeval now;
    get_current_time(&now);

    if (!timeout_head) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;  // No timeout, but select can handle NULL
        return;
    }

    *tv = timeout_head->expire_time;
    if (timeval_cmp(tv, &now) <= 0) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;
        return;
    }

    tv->tv_sec -= now.tv_sec;
    tv->tv_usec -= now.tv_usec;
    if (tv->tv_usec < 0) {
        tv->tv_sec--;
        tv->tv_usec += 1000000;
    }
}

void uasync_init(void) {
    FD_ZERO(&master_readfds);
    FD_ZERO(&master_writefds);
    FD_ZERO(&master_exceptfds);
    memset(fd_to_node, 0, sizeof(fd_to_node));  // Init map to NULL
}

void* uasync_set_timeout(int timeout, void* arg, timeout_callback_t callback) {
    if (timeout < 0 || !callback) return NULL;

    struct timeout_node* node = malloc(sizeof(struct timeout_node));
    if (!node) return NULL;

    get_current_time(&node->expire_time);
    timeval_add_tb(&node->expire_time, timeout);
    node->arg = arg;
    node->callback = callback;
    node->next = NULL;
    node->prev = NULL;

    return insert_timeout(node);
}

err_t uasync_cancel_timeout(void* t_id) {
    if (!t_id) return ERR_FAIL;

    struct timeout_node* node = (struct timeout_node*)t_id;
    if (node->prev) {
        node->prev->next = node->next;
    } else {
        timeout_head = node->next;
    }
    if (node->next) {
        node->next->prev = node->prev;
    }
    free(node);
    return ERR_OK;
}

void* uasync_add_socket(int fd, socket_callback_t read_cbk, socket_callback_t write_cbk, socket_callback_t except_cbk, void* user_data) {
    if (fd < 0 || fd >= FD_SETSIZE) return NULL;  // Add bounds check for map

    struct socket_node* node = malloc(sizeof(struct socket_node));
    if (!node) return NULL;

    node->fd = fd;
    node->read_cbk = read_cbk;
    node->write_cbk = write_cbk;
    node->except_cbk = except_cbk;
    node->user_data = user_data;
    node->next = socket_head;
    socket_head = node;

    // Update masters (point 1)
    if (read_cbk) FD_SET(fd, &master_readfds);
    if (write_cbk) FD_SET(fd, &master_writefds);
    if (except_cbk) FD_SET(fd, &master_exceptfds);

    // Update map (point 2)
    fd_to_node[fd] = node;

    if (fd > max_fd) max_fd = fd;

    return node;
}

err_t uasync_remove_socket(void* s_id) {
    if (!s_id) return ERR_FAIL;

    struct socket_node* node = (struct socket_node*)s_id;
    struct socket_node* cur = socket_head;
    struct socket_node* prev = NULL;

    while (cur) {
        if (cur == node) {
            if (prev) {
                prev->next = cur->next;
            } else {
                socket_head = cur->next;
            }

            // Update masters (point 1)
            if (node->read_cbk) FD_CLR(node->fd, &master_readfds);
            if (node->write_cbk) FD_CLR(node->fd, &master_writefds);
            if (node->except_cbk) FD_CLR(node->fd, &master_exceptfds);

            // Update map (point 2)
            fd_to_node[node->fd] = NULL;

            free(cur);

            // Update max_fd (simple rescan; optimize if needed by checking if removed == max_fd)
            max_fd = -1;
            cur = socket_head;
            while (cur) {
                if (cur->fd > max_fd) max_fd = cur->fd;
                cur = cur->next;
            }
            return ERR_OK;
        }
        prev = cur;
        cur = cur->next;
    }
    return ERR_FAIL;
}

void uasync_mainloop(void) {
    while (1) {
        uasync_poll(-1); /* infinite timeout */
    }
}

void uasync_poll(int timeout_tb) {
    /* Process expired timeouts */
    process_timeouts();

    /* Prepare select with copies of masters */
    fd_set readfds = master_readfds;
    fd_set writefds = master_writefds;
    fd_set exceptfds = master_exceptfds;

    struct timeval tv;
    get_next_timeout(&tv);
    
    /* If timeout_tb >= 0, compute timeout as min(timeout_tb, existing timer) */
    if (timeout_tb >= 0) {
        struct timeval user_tv;
        user_tv.tv_sec = timeout_tb / 10000;
        user_tv.tv_usec = (timeout_tb % 10000) * 100;
        
        /* If no internal timer or user timeout is smaller */
        if (tv.tv_sec == 0 && tv.tv_usec == 0 && !timeout_head) {
            tv = user_tv;
        } else if (user_tv.tv_sec < tv.tv_sec || 
                  (user_tv.tv_sec == tv.tv_sec && user_tv.tv_usec < tv.tv_usec)) {
            tv = user_tv;
        }
    }
    
    struct timeval* ptv = (tv.tv_sec == 0 && tv.tv_usec == 0 && !timeout_head) ? NULL : &tv;

    int nfds = select(max_fd + 1, &readfds, &writefds, &exceptfds, ptv);
    if (nfds < 0) {
        if (errno == EINTR) return;
        perror("select");
        return;
    }

    /* Process sockets with faster dispatch */
    for (int fd = 0; nfds > 0 && fd <= max_fd; fd++) {
        struct socket_node* node = fd_to_node[fd];
        if (!node) continue;

        if (node->except_cbk && FD_ISSET(fd, &exceptfds)) {
            node->except_cbk(fd, node->user_data);
            nfds--;
        }
        if (node->read_cbk && FD_ISSET(fd, &readfds)) {
            node->read_cbk(fd, node->user_data);
            nfds--;
        }
        if (node->write_cbk && FD_ISSET(fd, &writefds)) {
            node->write_cbk(fd, node->user_data);
            nfds--;
        }
    }
}