// uasync.c

#include "u_async.h"
#include "timeout_heap.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 {
    void* arg;
    timeout_callback_t callback;
    uint64_t expiration_ms; // absolute expiration time in milliseconds
};

// 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 TimeoutHeap* timeout_heap = NULL;  // Heap for timeout management
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 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;
}

// Convert timeval to milliseconds (uint64_t)
static uint64_t timeval_to_ms(const struct timeval* tv) {
    return (uint64_t)tv->tv_sec * 1000ULL + (uint64_t)tv->tv_usec / 1000ULL;
}



// Process expired timeouts
static void process_timeouts() {
    if (!timeout_heap) return;
    
    struct timeval now_tv;
    get_current_time(&now_tv);
    uint64_t now_ms = timeval_to_ms(&now_tv);
    
    while (1) {
        TimeoutEntry entry;
        if (timeout_heap_peek(timeout_heap, &entry) != 0) break;
        if (entry.expiration > now_ms) break;
        
        // Pop the expired timeout
        timeout_heap_pop(timeout_heap, &entry);
        struct timeout_node* node = (struct timeout_node*)entry.data;
        if (node && node->callback) {
            node->callback(node->arg);
        }
        free(node);
    }
}

// Compute time to next timeout
static void get_next_timeout(struct timeval* tv) {
    if (!timeout_heap) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;
        return;
    }
    
    TimeoutEntry entry;
    if (timeout_heap_peek(timeout_heap, &entry) != 0) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;
        return;
    }
    
    struct timeval now_tv;
    get_current_time(&now_tv);
    uint64_t now_ms = timeval_to_ms(&now_tv);
    
    if (entry.expiration <= now_ms) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;
        return;
    }
    
    uint64_t delta_ms = entry.expiration - now_ms;
    if (delta_ms > 86400000) { // Cap at 1 day to avoid overflow
        delta_ms = 86400000;
    }
    tv->tv_sec = delta_ms / 1000;
    tv->tv_usec = (delta_ms % 1000) * 1000;
}

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
    
    if (!timeout_heap) {
        timeout_heap = timeout_heap_create(16); // initial capacity 16
    }
}

void* uasync_set_timeout(int timeout_tb, void* arg, timeout_callback_t callback) {
    if (timeout_tb < 0 || !callback) return NULL;
    if (!timeout_heap) return NULL;
    
    struct timeout_node* node = malloc(sizeof(struct timeout_node));
    if (!node) return NULL;
    
    node->arg = arg;
    node->callback = callback;
    
    // Calculate expiration time in milliseconds
    struct timeval now;
    get_current_time(&now);
    timeval_add_tb(&now, timeout_tb);
    node->expiration_ms = timeval_to_ms(&now);
    
    // Insert into heap
    if (timeout_heap_push(timeout_heap, node->expiration_ms, node) != 0) {
        free(node);
        return NULL;
    }
    
    return node;
}

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

    struct timeout_node* node = (struct timeout_node*)t_id;
    
    // Try to cancel from heap
    if (timeout_heap_cancel(timeout_heap, node->expiration_ms, node) == 0) {
        free(node);
        return ERR_OK;
    }
    
    // If not found in heap (maybe already expired and removed), still free
    free(node);
    return ERR_FAIL;
}

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_heap || timeout_heap->size == 0)) {
            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_heap || timeout_heap->size == 0)) ? NULL : &tv;

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

    /* Process timeouts that may have expired during select */
    process_timeouts();

    /* 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--;
        }
    }
}