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

// Uasync instance structure
struct uasync_s {
    TimeoutHeap* timeout_heap;  // Heap for timeout management
    struct socket_node* socket_head;
    int max_fd;
    fd_set master_readfds;
    fd_set master_writefds;
    fd_set master_exceptfds;
    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;
}

// Instance management
uasync_t* uasync_create(void) {
    uasync_t* ua = malloc(sizeof(uasync_t));
    if (!ua) return NULL;
    
    memset(ua, 0, sizeof(uasync_t));
    ua->max_fd = -1;
    FD_ZERO(&ua->master_readfds);
    FD_ZERO(&ua->master_writefds);
    FD_ZERO(&ua->master_exceptfds);
    memset(ua->fd_to_node, 0, sizeof(ua->fd_to_node));
    
    ua->timeout_heap = timeout_heap_create(16);
    if (!ua->timeout_heap) {
        free(ua);
        return NULL;
    }
    
    return ua;
}

void uasync_destroy(uasync_t* ua) {
    if (!ua) return;
    
    // Free all remaining timeouts
    if (ua->timeout_heap) {
        while (1) {
            TimeoutEntry entry;
            if (timeout_heap_pop(ua->timeout_heap, &entry) != 0) break;
            struct timeout_node* node = (struct timeout_node*)entry.data;
            free(node);
        }
        timeout_heap_destroy(ua->timeout_heap);
    }
    
    // Free all socket nodes
    struct socket_node* cur = ua->socket_head;
    while (cur) {
        struct socket_node* next = cur->next;
        free(cur);
        cur = next;
    }
    
    free(ua);
}

void uasync_init(uasync_t* ua) {
    if (!ua) return;
    
    ua->max_fd = -1;
    FD_ZERO(&ua->master_readfds);
    FD_ZERO(&ua->master_writefds);
    FD_ZERO(&ua->master_exceptfds);
    memset(ua->fd_to_node, 0, sizeof(ua->fd_to_node));
    
    if (!ua->timeout_heap) {
        ua->timeout_heap = timeout_heap_create(16);
    }
}

// Process expired timeouts
static void process_timeouts(uasync_t* ua) {
    if (!ua || !ua->timeout_heap) return;
    
    struct timeval now_tv;
    get_current_time(&now_tv);
    uint64_t now_ms = timeval_to_ms(&now_tv);
    
    int processed = 0;
    while (1) {
        TimeoutEntry entry;
        if (timeout_heap_peek(ua->timeout_heap, &entry) != 0) break;
        if (entry.expiration > now_ms) break;
        
        // DEBUG
        // printf("process_timeouts: expiration=%lu now=%lu diff=%ld heap_size=%zu\n", 
        //        entry.expiration, now_ms, (long)(now_ms - entry.expiration),
        //        ua->timeout_heap->size);
        
        // Pop the expired timeout
        timeout_heap_pop(ua->timeout_heap, &entry);
        struct timeout_node* node = (struct timeout_node*)entry.data;
        if (node && node->callback) {
            node->callback(node->arg);
            processed++;
        } else {
            printf("WARN: timeout node missing callback\n");
        }
        free(node);
    }
    if (processed > 0) {
        // printf("process_timeouts: processed %d timeouts, now_ms=%lu heap_size=%zu\n", 
        //        processed, now_ms, ua->timeout_heap->size);
    }
}

// Compute time to next timeout
static void get_next_timeout(uasync_t* ua, struct timeval* tv) {
    if (!ua || !ua->timeout_heap) {
        tv->tv_sec = 0;
        tv->tv_usec = 0;
        return;
    }
    
    TimeoutEntry entry;
    if (timeout_heap_peek(ua->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;
    
    // DEBUG
    // printf("get_next_timeout: expiration=%lu now=%lu delta=%lu sec=%ld usec=%ld\n",
    //        entry.expiration, now_ms, delta_ms, (long)tv->tv_sec, (long)tv->tv_usec);
}



void* uasync_set_timeout(uasync_t* ua, int timeout_tb, void* arg, timeout_callback_t callback) {
    // printf("SETTIMEOUT called\n");
    // fflush(stdout);
    if (!ua || timeout_tb < 0 || !callback) return NULL;
    if (!ua->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);
    uint64_t now_ms = timeval_to_ms(&now);
    timeval_add_tb(&now, timeout_tb);
    node->expiration_ms = timeval_to_ms(&now);
    
    // DEBUG
    // printf("set_timeout: tb=%d (%.2fms), now=%lu, expiration=%lu, delta=%ldms\n", 
    //        timeout_tb, timeout_tb/10.0, now_ms, node->expiration_ms, 
    //        (long)(node->expiration_ms - now_ms));
    
    // Insert into heap
    if (timeout_heap_push(ua->timeout_heap, node->expiration_ms, node) != 0) {
        free(node);
        return NULL;
    }
    
    return node;
}



err_t uasync_cancel_timeout(uasync_t* ua, void* t_id) {
    if (!ua || !t_id || !ua->timeout_heap) return ERR_FAIL;

    struct timeout_node* node = (struct timeout_node*)t_id;
    
    // Try to cancel from heap
    if (timeout_heap_cancel(ua->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(uasync_t* ua, int fd, socket_callback_t read_cbk, socket_callback_t write_cbk, socket_callback_t except_cbk, void* user_data) {
    if (!ua || 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 = ua->socket_head;
    ua->socket_head = node;

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

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

    if (fd > ua->max_fd) ua->max_fd = fd;

    return node;
}



err_t uasync_remove_socket(uasync_t* ua, void* s_id) {
    if (!ua || !s_id) return ERR_FAIL;

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

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

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

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

            free(cur);

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



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



void uasync_poll(uasync_t* ua, int timeout_tb) {
    if (!ua) return;
    
    /* Process expired timeouts */
    process_timeouts(ua);

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

    struct timeval tv;
    get_next_timeout(ua, &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 && (!ua->timeout_heap || ua->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 && (!ua->timeout_heap || ua->timeout_heap->size == 0)) ? NULL : &tv;

    int nfds = select(ua->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 <= ua->max_fd; fd++) {
        struct socket_node* node = ua->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--;
        }
    }

    /* Process timeouts that may have expired during select. тайм-ауты после данных. */
    process_timeouts(ua);

}