#include <c_TcpPollReactor.h>

#if !defined(_WIN32) && !defined(_WIN64)
    #include <poll.h>
#endif

typedef struct {
    c_TcpReactor_t base;
    struct pollfd* fds;
    c_size_t       count;
    c_size_t       capacity;
    c_Allocator_t  allocator;
} c_PollReactorImpl_t;

/* (Adheres to the O(1) packing allocations logic established previously) */
static c_err_t Poll_Add(c_TcpReactor_t* self, c_socket_t sock, uint32_t events) {
    c_PollReactorImpl_t* impl = (c_PollReactorImpl_t*)self;
    if (impl->count >= impl->capacity) {
        c_size_t old_cap = impl->capacity;
        c_size_t new_cap = old_cap == 0 ? 16 : old_cap * 2;
        struct pollfd* new_fds = (struct pollfd*)c_Allocator_Realloc(&impl->allocator, impl->fds, old_cap * sizeof(struct pollfd), new_cap * sizeof(struct pollfd));
        if (!new_fds) return C_ERR_NOMEM;
        impl->fds = new_fds;
        impl->capacity = new_cap;
    }
    impl->fds[impl->count].fd = (int)sock;
    impl->fds[impl->count].events = (short)(events ? POLLIN : 0);
    impl->fds[impl->count].revents = 0;
    impl->count++;
    return C_SUCCESS;
}

static c_err_t Poll_Remove(c_TcpReactor_t* self, c_socket_t sock) {
    c_PollReactorImpl_t* impl = (c_PollReactorImpl_t*)self;
    for (c_size_t i = 0; i < impl->count; ++i) {
        if (impl->fds[i].fd == (int)sock) {
            impl->fds[i] = impl->fds[impl->count - 1]; /* O(1) tail-swap */
            impl->count--;
            return C_SUCCESS;
        }
    }
    return C_ERR_NOTFOUND;
}

static c_err_t Poll_Poll(c_TcpReactor_t* self, long timeout_ms, c_TcpServer_t* server, void* args) {
    c_PollReactorImpl_t* impl = (c_PollReactorImpl_t*)self;
    if (impl->count == 0 || !server->is_running) return C_SUCCESS;

#if defined(_WIN32) || defined(_WIN64)
    int ret = WSAPoll(impl->fds, (ULONG)impl->count, timeout_ms);
#else
    int ret = poll(impl->fds, (nfds_t)impl->count, timeout_ms);
#endif

    if (ret < 0) {
        if (server->fnOnError) {
            server->fnOnError(server, server->listen_sock, C_TCPSERVER_ERR_ON_POLL, args);
        }
    }
    if (ret == 0) return C_SUCCESS;

    for (c_size_t i = impl->count; i > 0; --i) {
        c_size_t idx = i - 1;
        struct pollfd* current_fd = &impl->fds[idx];

        if (current_fd->revents == 0) continue;

        if (current_fd->fd == (int)server->listen_sock) {
            if (current_fd->revents & POLLIN) {
                c_SockAddr_t peer_addr;
                c_socket_t client = c_SocketUtil_Accept(server->listen_sock, &peer_addr);

                if (c_Socket_IsValid(client)) {
                    bool keep = true;
                    if (server->fnOnConnect) {
                        keep = server->fnOnConnect(server, client, &peer_addr, args);
                    }

                    if (keep) {
                        Poll_Add(self, client, 1);
                    } else {
                        c_Socket_Close(client);
                    }
                }else if (server->fnOnError) {
                    server->fnOnError(server, server->listen_sock, C_TCPSERVER_ERR_ON_ACCEPT, args);
                }
            }else if (current_fd->revents & (POLLERR | POLLHUP)) {
                if (server->fnOnError) {
                    server->fnOnError(server, server->listen_sock, C_TCPSERVER_ERR_ON_POLLERR, args);
                }
            }
            ret--;
        }
        else {
            c_socket_t client_sock = (c_socket_t)current_fd->fd;

            if (current_fd->revents & (POLLIN | POLLHUP | POLLERR)) {
                c_bool_t keep_alive = C_TRUE;

                /* Trigger error logging if a structural socket error bit is present */
                if (current_fd->revents & POLLERR) {
                    if (server->fnOnError) {
                        server->fnOnError(server, client_sock, C_TCPSERVER_ERR_ON_POLLERR, args);
                    }
                    keep_alive = C_FALSE;
                }

                if (keep_alive) {
                    char peek_buf;
#if defined(_WIN32) || defined(_WIN64)
                    int peek_res = recv(client_sock, &peek_buf, 1, MSG_PEEK);
#else
                    ssize_t peek_res = recv(client_sock, &peek_buf, 1, MSG_PEEK);
#endif

                    if (peek_res == 0 || peek_res == C_SOCKET_ERROR) {
                        keep_alive = C_FALSE;
                    } else if (server->fnOnRequest) {
                        keep_alive = server->fnOnRequest(server, client_sock, args);
                    }
                }

                if (!keep_alive) {
                    Poll_Remove(self, client_sock);

                    /* Fire the explicitly realigned c_TcpServer_OnDisconnect_t call hook */
                    if (server->fnOnDisconnect) {
                        server->fnOnDisconnect(server, client_sock, args);
                    }

                    c_Socket_Close(client_sock);
                }
            }
            ret--;
        }

        if (ret == 0) break;
    }
    return C_SUCCESS;
}

/* (The rest of the driver functions stay identically compiled) */
static void Poll_Destroy(c_TcpReactor_t* self) {
    c_PollReactorImpl_t* impl = (c_PollReactorImpl_t*)self;
    c_Allocator_t alloc = impl->allocator;
    for (c_size_t i = 0; i < impl->count; ++i) {
        c_socket_t sock = (c_socket_t)impl->fds[i].fd;
        c_Socket_Close(sock);
    }
    if (impl->fds) c_Allocator_Free(&alloc, impl->fds);
    c_Allocator_Free(&alloc, impl);
}

static const c_TcpReactorVtbl_t g_PollReactorVtbl = { Poll_Add, Poll_Remove, Poll_Poll, Poll_Destroy };

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

c_err_t c_TcpPollReactor_Create(c_TcpReactor_t** out_reactor, c_Allocator_t* allocator) {
    if (!out_reactor) return C_ERR_PARAM;
    c_Allocator_t alloc = allocator ? *allocator : c_DefaultAllocator;
    c_PollReactorImpl_t* impl = (c_PollReactorImpl_t*)c_Allocator_Calloc(&alloc, 1, sizeof(c_PollReactorImpl_t));
    if (!impl) return C_ERR_NOMEM;
    impl->base.vtbl = &g_PollReactorVtbl;
    impl->allocator = alloc;
    *out_reactor = &impl->base;
    return C_SUCCESS;
}