#include "pxl8_net.h"

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

#ifdef _WIN32
    #define WIN32_LEAN_AND_MEAN
    #include <winsock2.h>
    #include <ws2tcpip.h>
    typedef SOCKET socket_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
#define PXL8_NET_TICK_RATE 30.0f

struct pxl8_net {
    char address[256];
    bool connected;
    pxl8_entity_state entities[PXL8_MAX_SNAPSHOT_ENTITIES];
    pxl8_event_msg events[PXL8_MAX_SNAPSHOT_EVENTS];
    u64 highest_tick;
    pxl8_input_msg input_history[PXL8_NET_INPUT_HISTORY_SIZE];
    u64 input_head;
    u64 input_oldest_tick;
    f32 interp_time;
    pxl8_net_mode mode;
    u16 port;
    u8 predicted_state[PXL8_NET_USERDATA_SIZE];
    u64 predicted_tick;
    pxl8_entity_state prev_entities[PXL8_MAX_SNAPSHOT_ENTITIES];
    pxl8_snapshot_header prev_snapshot;
    u8 recv_buf[4096];
    u8 send_buf[4096];
    u32 sequence;
    struct sockaddr_in server_addr;
    pxl8_snapshot_header snapshot;
    socket_t sock;
};

static const pxl8_entity_state* find_entity(const pxl8_entity_state* entities, u16 count, u64 id) {
    for (u16 i = 0; i < count; i++) {
        if (entities[i].entity_id == id) return &entities[i];
    }
    return NULL;
}

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

    memset(&net->server_addr, 0, sizeof(net->server_addr));
    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 = calloc(1, sizeof(pxl8_net));
    if (!net) return NULL;

    net->mode = config->mode;
    net->port = config->port ? config->port : PXL8_NET_DEFAULT_PORT;
    net->sock = INVALID_SOCK;
    net->connected = false;
    net->sequence = 0;
    net->highest_tick = 0;

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

const pxl8_entity_state* pxl8_net_entities(const pxl8_net* net) {
    if (!net) return NULL;
    return net->entities;
}

const u8* pxl8_net_entity_prev_userdata(const pxl8_net* net, u64 entity_id) {
    if (!net) return NULL;
    const pxl8_entity_state* e = find_entity(net->prev_entities, net->prev_snapshot.entity_count, entity_id);
    return e ? e->userdata : NULL;
}

const u8* pxl8_net_entity_userdata(const pxl8_net* net, u64 entity_id) {
    if (!net) return NULL;
    const pxl8_entity_state* e = find_entity(net->entities, net->snapshot.entity_count, entity_id);
    return e ? e->userdata : NULL;
}

const pxl8_event_msg* pxl8_net_events(const pxl8_net* net) {
    if (!net) return NULL;
    return net->events;
}

const pxl8_input_msg* pxl8_net_input_at(const pxl8_net* net, u64 tick) {
    if (!net) return NULL;
    for (u64 i = 0; i < PXL8_NET_INPUT_HISTORY_SIZE; i++) {
        if (net->input_history[i].tick == tick) {
            return &net->input_history[i];
        }
    }
    return NULL;
}

u64 pxl8_net_input_oldest_tick(const pxl8_net* net) {
    if (!net) return 0;
    return net->input_oldest_tick;
}

void pxl8_net_input_push(pxl8_net* net, const pxl8_input_msg* input) {
    if (!net || !input) return;
    u64 idx = net->input_head % PXL8_NET_INPUT_HISTORY_SIZE;
    net->input_history[idx] = *input;
    net->input_head++;
    if (net->input_oldest_tick == 0 || input->tick < net->input_oldest_tick) {
        net->input_oldest_tick = input->tick;
    }
}

f32 pxl8_net_lerp_alpha(const pxl8_net* net) {
    if (!net) return 1.0f;
    f32 tick_duration = 1.0f / PXL8_NET_TICK_RATE;
    f32 alpha = net->interp_time / tick_duration;
    return alpha > 1.0f ? 1.0f : alpha;
}

bool pxl8_net_needs_correction(const pxl8_net* net) {
    if (!net) return false;
    if (net->snapshot.tick == 0) return false;
    if (net->predicted_tick == 0) return false;
    if (net->snapshot.tick > net->predicted_tick) return true;
    const u8* server = pxl8_net_entity_userdata(net, net->snapshot.player_id);
    if (!server) return false;
    return memcmp(server, net->predicted_state, PXL8_NET_USERDATA_SIZE) != 0;
}

u64 pxl8_net_player_id(const pxl8_net* net) {
    if (!net) return 0;
    return net->snapshot.player_id;
}

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

    size_t len = pxl8_net_recv(net, net->recv_buf, sizeof(net->recv_buf));
    if (len < sizeof(pxl8_msg_header)) return false;

    pxl8_msg_header hdr;
    size_t offset = pxl8_protocol_deserialize_header(net->recv_buf, len, &hdr);
    if (hdr.type != PXL8_MSG_SNAPSHOT) return false;

    pxl8_snapshot_header snap;
    offset += pxl8_protocol_deserialize_snapshot_header(net->recv_buf + offset, len - offset, &snap);

    if (snap.tick <= net->highest_tick) return false;

    memcpy(net->prev_entities, net->entities, sizeof(net->entities));
    net->prev_snapshot = net->snapshot;

    net->highest_tick = snap.tick;
    net->snapshot = snap;
    net->interp_time = 0.0f;

    u16 count = snap.entity_count;
    if (count > PXL8_MAX_SNAPSHOT_ENTITIES) count = PXL8_MAX_SNAPSHOT_ENTITIES;

    for (u16 i = 0; i < count; i++) {
        offset += pxl8_protocol_deserialize_entity_state(
            net->recv_buf + offset, len - offset, &net->entities[i]);
    }

    return true;
}

u8* pxl8_net_predicted_state(pxl8_net* net) {
    if (!net) return NULL;
    return net->predicted_state;
}

void pxl8_net_predicted_tick_set(pxl8_net* net, u64 tick) {
    if (!net) return;
    net->predicted_tick = tick;
}

size_t pxl8_net_recv(pxl8_net* net, u8* buf, size_t 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) ? (size_t)received : 0;
}

pxl8_result pxl8_net_send(pxl8_net* net, const u8* data, size_t 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;
}

pxl8_result pxl8_net_send_command(pxl8_net* net, const pxl8_command_msg* cmd) {
    if (!net || !net->connected) return PXL8_ERROR_INVALID_ARGUMENT;

    u8 buf[sizeof(pxl8_msg_header) + sizeof(pxl8_command_msg)];
    pxl8_msg_header hdr = {
        .type = PXL8_MSG_COMMAND,
        .version = PXL8_PROTOCOL_VERSION,
        .size = sizeof(pxl8_command_msg),
        .sequence = 0
    };

    size_t offset = pxl8_protocol_serialize_header(&hdr, buf, sizeof(buf));
    offset += pxl8_protocol_serialize_command(cmd, buf + offset, sizeof(buf) - offset);

    return pxl8_net_send(net, buf, offset);
}

pxl8_result pxl8_net_send_input(pxl8_net* net, const pxl8_input_msg* input) {
    if (!net || !net->connected) return PXL8_ERROR_INVALID_ARGUMENT;

    u8 buf[sizeof(pxl8_msg_header) + sizeof(pxl8_input_msg)];
    pxl8_msg_header hdr = {
        .type = PXL8_MSG_INPUT,
        .version = PXL8_PROTOCOL_VERSION,
        .size = sizeof(pxl8_input_msg),
        .sequence = 0
    };

    size_t offset = pxl8_protocol_serialize_header(&hdr, buf, sizeof(buf));
    offset += pxl8_protocol_serialize_input(input, buf + offset, sizeof(buf) - offset);

    return pxl8_net_send(net, buf, offset);
}

const pxl8_snapshot_header* pxl8_net_snapshot(const pxl8_net* net) {
    if (!net) return NULL;
    return &net->snapshot;
}

u64 pxl8_net_tick(const pxl8_net* net) {
    if (!net) return 0;
    return net->snapshot.tick;
}

void pxl8_net_update(pxl8_net* net, f32 dt) {
    if (!net) return;
    net->interp_time += dt;
}