#include "pxl8_net.h"

#include <stdlib.h>
#include <string.h>

#include "pxl8_platform.h"

#ifdef PXL8_ASYNC_THREADS
#include <stdatomic.h>
#include "pxl8_queue.h"
#endif

#include "pxl8_log.h"
#include "pxl8_mem.h"

#ifdef _WIN32
    #define WIN32_LEAN_AND_MEAN
    #include <winsock2.h>
    #include <ws2tcpip.h>
    typedef SOCKET socket_t;
    typedef int ssize_t;
    #define INVALID_SOCK INVALID_SOCKET
    #define close_socket closesocket
#else
    #include <arpa/inet.h>
    #include <errno.h>
    #include <fcntl.h>
    #include <netinet/in.h>
    #include <sys/socket.h>
    #include <unistd.h>
    typedef int socket_t;
    #define INVALID_SOCK -1
    #define close_socket close
#endif

#define PXL8_NET_DEFAULT_PORT 7777

struct pxl8_net {
    char address[256];
    bool connected;
    u16 port;
    struct sockaddr_in server_addr;
    socket_t sock;

    u8 recv_buf[4096];
    u8 send_buf[4096];

#ifdef PXL8_ASYNC_THREADS
    pxl8_thread* recv_thread;
    atomic_bool running;
    pxl8_queue recv_queue;
#endif
};

pxl8_result pxl8_net_connect(pxl8_net* net) {
    if (!net) return PXL8_ERROR_INVALID_ARGUMENT;
    if (net->connected) return PXL8_OK;

#ifdef _WIN32
    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) {
        return PXL8_ERROR_SYSTEM_FAILURE;
    }
#endif

    net->sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (net->sock == INVALID_SOCK) {
        return PXL8_ERROR_SYSTEM_FAILURE;
    }

#ifdef _WIN32
    u_long nonblocking = 1;
    ioctlsocket(net->sock, FIONBIO, &nonblocking);
#else
    int flags = fcntl(net->sock, F_GETFL, 0);
    fcntl(net->sock, F_SETFL, flags | O_NONBLOCK);
#endif

    int rcvbuf = 1024 * 1024;
    setsockopt(net->sock, SOL_SOCKET, SO_RCVBUF, (const char*)&rcvbuf, sizeof(rcvbuf));

    memset(&net->server_addr, 0, sizeof(net->server_addr));
#ifdef __APPLE__
    net->server_addr.sin_len = sizeof(net->server_addr);
#endif
    net->server_addr.sin_family = AF_INET;
    net->server_addr.sin_port = htons(net->port);
    inet_pton(AF_INET, net->address, &net->server_addr.sin_addr);

    net->connected = true;
    return PXL8_OK;
}

bool pxl8_net_connected(const pxl8_net* net) {
    return net && net->connected;
}

pxl8_net* pxl8_net_create(const pxl8_net_config* config) {
    pxl8_net* net = pxl8_calloc(1, sizeof(pxl8_net));
    if (!net) return NULL;

    net->port = config->port ? config->port : PXL8_NET_DEFAULT_PORT;
    net->sock = INVALID_SOCK;

    if (config->address) {
        strncpy(net->address, config->address, sizeof(net->address) - 1);
    } else {
        strncpy(net->address, "127.0.0.1", sizeof(net->address) - 1);
    }

    return net;
}

void pxl8_net_destroy(pxl8_net* net) {
    if (!net) return;
    pxl8_net_disconnect(net);
    pxl8_free(net);
}

void pxl8_net_disconnect(pxl8_net* net) {
    if (!net) return;
    if (net->sock != INVALID_SOCK) {
        close_socket(net->sock);
        net->sock = INVALID_SOCK;
    }
#ifdef _WIN32
    WSACleanup();
#endif
    net->connected = false;
}

bool pxl8_net_poll(pxl8_net* net) {
    if (!net || !net->connected) return false;

    usize len = pxl8_net_recv(net, net->recv_buf, sizeof(net->recv_buf));
    return len > 0;
}

usize pxl8_net_recv(pxl8_net* net, u8* buf, usize len) {
    if (!net || !net->connected) return 0;

    struct sockaddr_in from;
    socklen_t from_len = sizeof(from);
    ssize_t received = recvfrom(net->sock, (char*)buf, len, 0,
                                 (struct sockaddr*)&from, &from_len);
    return (received > 0) ? (usize)received : 0;
}

pxl8_result pxl8_net_send(pxl8_net* net, const u8* data, usize len) {
    if (!net || !net->connected) return PXL8_ERROR_INVALID_ARGUMENT;

    ssize_t sent = sendto(net->sock, (const char*)data, len, 0,
                          (struct sockaddr*)&net->server_addr,
                          sizeof(net->server_addr));
    return (sent > 0) ? PXL8_OK : PXL8_ERROR_SYSTEM_FAILURE;
}

#ifdef PXL8_ASYNC_THREADS

static int pxl8_net_recv_thread(void* data) {
    pxl8_net* net = (pxl8_net*)data;
    u8 buf[PXL8_NET_PACKET_MAX_SIZE];

    while (atomic_load(&net->running)) {
        if (!net->connected) {
            pxl8_sleep_ms(10);
            continue;
        }

        struct sockaddr_in from;
        socklen_t from_len = sizeof(from);
        ssize_t received = recvfrom(net->sock, (char*)buf, sizeof(buf), 0,
                                     (struct sockaddr*)&from, &from_len);

        if (received > 0) {
            pxl8_packet* pkt = pxl8_malloc(sizeof(pxl8_packet));
            if (pkt) {
                memcpy(pkt->data, buf, (usize)received);
                pkt->len = (usize)received;

                if (!pxl8_queue_push(&net->recv_queue, pkt)) {
                    pxl8_free(pkt);
                }
            }
        } else {
            pxl8_sleep_ms(1);
        }
    }

    return 0;
}

void pxl8_net_start_thread(pxl8_net* net) {
    if (!net || net->recv_thread) return;

    pxl8_queue_init(&net->recv_queue);
    atomic_store(&net->running, true);
    net->recv_thread = pxl8_thread_create(pxl8_net_recv_thread, "pxl8_net_recv", net);
}

void pxl8_net_stop_thread(pxl8_net* net) {
    if (!net || !net->recv_thread) return;

    atomic_store(&net->running, false);
    pxl8_thread_wait(net->recv_thread, NULL);
    net->recv_thread = NULL;

    pxl8_packet* pkt;
    while ((pkt = pxl8_queue_pop(&net->recv_queue))) {
        pxl8_free(pkt);
    }
}

pxl8_packet* pxl8_net_pop_packet(pxl8_net* net) {
    if (!net) return NULL;
    return (pxl8_packet*)pxl8_queue_pop(&net->recv_queue);
}

void pxl8_net_packet_free(pxl8_packet* pkt) {
    pxl8_free(pkt);
}

#endif