#include <c_TcpClientPollReactor.h>

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

typedef struct {
    c_TcpClientReactor_t base;
    struct pollfd*       fds;
    c_size_t             count;
    c_size_t             capacity;
    c_Allocator_t        allocator;
} c_ClientPollImpl_t;

static c_err_t ClientPoll_Add(c_TcpClientReactor_t* self, c_socket_t sock, uint32_t events) {
    c_ClientPollImpl_t* impl = (c_ClientPollImpl_t*)self;
    if (impl->count >= impl->capacity) {
        c_size_t old_cap = impl->capacity;
        c_size_t new_cap = old_cap == 0 ? 4 : 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 ClientPoll_Remove(c_TcpClientReactor_t* self, c_socket_t sock) {
    c_ClientPollImpl_t* impl = (c_ClientPollImpl_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 ClientPoll_Poll(c_TcpClientReactor_t* self, long timeout_ms, c_TcpClient_t* client, void* args) {
    c_ClientPollImpl_t* impl = (c_ClientPollImpl_t*)self;
    if (!client->is_linked || !c_Socket_IsValid(client->sock)) return C_SUCCESS;

    if (impl->count == 0) {
        ClientPoll_Add(self, client->sock, 1);
    }

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

    if (ret < 0) {
        if (client->fnOnError) {
            client->fnOnError(client, C_TCPCLIENT_ERR_ON_POLL, args);
        }
        return C_SUCCESS;
    }
    if (ret == 0) return C_SUCCESS;

    struct pollfd* client_fd = &impl->fds[0];
    if (client_fd->revents & (POLLIN | POLLHUP | POLLERR)) {
        c_bool_t keep_alive = C_TRUE;

        if (client_fd->revents & POLLERR) {
            if (client->fnOnError) {
                client->fnOnError(client, C_TCPCLIENT_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 (client->fnOnData) {
                keep_alive = client->fnOnData(client, args);
            }
        }

        if (!keep_alive || (client_fd->revents & POLLHUP)) {
            ClientPoll_Remove(self, client->sock);
            if (client->fnOnDisconnect) {
                client->fnOnDisconnect(client, args);
            }
            c_Socket_Close(client->sock);
            client->is_linked = false;
            return C_ERR_EOF;
        }
    }
    return C_SUCCESS;
}

static void ClientPoll_Destroy(c_TcpClientReactor_t* self) {
    c_ClientPollImpl_t* impl = (c_ClientPollImpl_t*)self;
    c_Allocator_t alloc = impl->allocator;
    if (impl->fds) c_Allocator_Free(&alloc, impl->fds);
    c_Allocator_Free(&alloc, impl);
}

static const c_TcpClientReactorVtbl_t g_ClientPollVtbl = { ClientPoll_Add, ClientPoll_Remove, ClientPoll_Poll, ClientPoll_Destroy };

c_err_t c_TcpClientPollReactor_Create(c_TcpClientReactor_t** out_reactor, c_Allocator_t* allocator) {
    if (!out_reactor) return C_ERR_PARAM;
    c_Allocator_t alloc = allocator ? *allocator : c_DefaultAllocator;
    c_ClientPollImpl_t* impl = (c_ClientPollImpl_t*)c_Allocator_Calloc(&alloc, 1, sizeof(c_ClientPollImpl_t));
    if (!impl) return C_ERR_NOMEM;
    impl->base.vtbl = &g_ClientPollVtbl;
    impl->fds = NULL;
    impl->count = 0;
    impl->capacity = 0;
    impl->allocator = alloc;
    *out_reactor = &impl->base;
    return C_SUCCESS;
}