#include "pxl8_net.h"

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

#include "pxl8_hal.h"

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

#include "pxl8_bytes.h"
#include "pxl8_log.h"
#include "pxl8_mem.h"
#include "pxl8_world.h"
#include "pxl8_world_chunk_cache.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;
    u16 port;
    struct sockaddr_in server_addr;
    socket_t sock;

    pxl8_world_chunk_cache* chunk_cache;
    u32 chunk_id;
    u8 chunk_type;
    pxl8_world* world;

    f32 dt;
    u64 highest_tick;
    f32 interp_time;
    u32 sequence;

    pxl8_entity_state entities[PXL8_MAX_SNAPSHOT_ENTITIES];
    pxl8_event_msg events[PXL8_MAX_SNAPSHOT_EVENTS];
    pxl8_entity_state prev_entities[PXL8_MAX_SNAPSHOT_ENTITIES];
    pxl8_snapshot_header prev_snapshot;
    pxl8_snapshot_header snapshot;

    u64 input_head;
    pxl8_input_msg input_history[PXL8_NET_INPUT_HISTORY_SIZE];
    u64 input_oldest_tick;

    u8 predicted_state[PXL8_NET_USERDATA_SIZE];
    u64 predicted_tick;

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

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

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 = pxl8_calloc(1, sizeof(pxl8_net));
    if (!net) return NULL;

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

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;

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

    pxl8_msg_header hdr;
    usize offset = pxl8_protocol_deserialize_header(net->recv_buf, len, &hdr);

    if (hdr.type == PXL8_MSG_CHUNK) {
        if (!net->chunk_cache) return false;

        pxl8_chunk_msg_header chunk_hdr;
        offset += pxl8_protocol_deserialize_chunk_msg_header(net->recv_buf + offset, len - offset, &chunk_hdr);

        const u8* payload = net->recv_buf + offset;
        usize payload_len = chunk_hdr.payload_size;
        if (payload_len > len - offset) {
            payload_len = len - offset;
        }

        pxl8_world_chunk_cache_receive(net->chunk_cache, &chunk_hdr, payload, payload_len);
        return true;
    }

    if (hdr.type == PXL8_MSG_CHUNK_ENTER) {
        pxl8_chunk_enter_msg chunk_msg;
        pxl8_protocol_deserialize_chunk_enter(net->recv_buf + offset, len - offset, &chunk_msg);
        net->chunk_id = chunk_msg.chunk_id;
        net->chunk_type = chunk_msg.chunk_type;
        pxl8_debug("[CLIENT] Received CHUNK_ENTER type=%u id=%u", chunk_msg.chunk_type, chunk_msg.chunk_id);
        return true;
    }

    if (hdr.type == PXL8_MSG_CHUNK_EXIT) {
        net->chunk_id = 0;
        net->chunk_type = 0;
        return true;
    }

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

    if (net->world) {
        pxl8_world_reconcile(net->world, net, 1.0f / 30.0f);
    }

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

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

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

    usize 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;

    pxl8_net_input_push(net, input);

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

    usize 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->dt = dt;
    net->interp_time += dt;
}

void pxl8_net_set_chunk_cache(pxl8_net* net, pxl8_world_chunk_cache* cache) {
    if (!net) return;
    net->chunk_cache = cache;
}

pxl8_world_chunk_cache* pxl8_net_chunk_cache(pxl8_net* net) {
    if (!net) return NULL;
    return net->chunk_cache;
}

void pxl8_net_set_world(pxl8_net* net, pxl8_world* world) {
    if (!net) return;
    net->world = world;
}

u32 pxl8_net_chunk_id(const pxl8_net* net) {
    if (!net) return 0;
    return net->chunk_id;
}

u8 pxl8_net_chunk_type(const pxl8_net* net) {
    if (!net) return PXL8_CHUNK_TYPE_VXL;
    return net->chunk_type;
}

pxl8_result pxl8_net_send_chunk_settings(pxl8_net* net, i32 render_distance, i32 sim_distance) {
    if (!net) return PXL8_ERROR_NULL_POINTER;
    if (!net->connected) return PXL8_ERROR_NOT_CONNECTED;

    pxl8_command_msg cmd = {0};
    cmd.cmd_type = PXL8_CMD_SET_CHUNK_SETTINGS;
    pxl8_pack_i32_be(cmd.payload, 0, render_distance);
    pxl8_pack_i32_be(cmd.payload, 4, sim_distance);
    cmd.payload_size = 8;

    return pxl8_net_send_command(net, &cmd);
}

pxl8_result pxl8_net_spawn(pxl8_net* net, f32 x, f32 y, f32 z, f32 yaw, f32 pitch) {
    if (!net) return PXL8_ERROR_NULL_POINTER;
    if (!net->connected) return PXL8_ERROR_NOT_CONNECTED;

    pxl8_command_msg cmd = {0};
    cmd.cmd_type = PXL8_CMD_SPAWN_ENTITY;
    pxl8_pack_f32_be(cmd.payload, 0, x);
    pxl8_pack_f32_be(cmd.payload, 4, y);
    pxl8_pack_f32_be(cmd.payload, 8, z);
    pxl8_pack_f32_be(cmd.payload, 12, yaw);
    pxl8_pack_f32_be(cmd.payload, 16, pitch);
    cmd.payload_size = 20;

    return pxl8_net_send_command(net, &cmd);
}

pxl8_result pxl8_net_exit_chunk(pxl8_net* net, f32 x, f32 y, f32 z) {
    if (!net) return PXL8_ERROR_NULL_POINTER;
    if (!net->connected) return PXL8_ERROR_NOT_CONNECTED;

    pxl8_command_msg cmd = {0};
    cmd.cmd_type = PXL8_CMD_EXIT_CHUNK;
    pxl8_pack_f32_be(cmd.payload, 0, x);
    pxl8_pack_f32_be(cmd.payload, 4, y);
    pxl8_pack_f32_be(cmd.payload, 8, z);
    cmd.payload_size = 12;

    return pxl8_net_send_command(net, &cmd);
}

pxl8_result pxl8_net_enter_chunk(pxl8_net* net, u32 chunk_id) {
    if (!net) return PXL8_ERROR_NULL_POINTER;
    if (!net->connected) return PXL8_ERROR_NOT_CONNECTED;

    pxl8_command_msg cmd = {0};
    cmd.cmd_type = PXL8_CMD_ENTER_CHUNK;
    pxl8_pack_u32_be(cmd.payload, 0, chunk_id);
    cmd.payload_size = 4;

    return pxl8_net_send_command(net, &cmd);
}

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

bool pxl8_net_process_packet(pxl8_net* net, const pxl8_packet* pkt) {
    if (!net || !pkt || pkt->len < sizeof(pxl8_msg_header)) return false;

    pxl8_msg_header hdr;
    usize offset = pxl8_protocol_deserialize_header(pkt->data, pkt->len, &hdr);

    if (hdr.type == PXL8_MSG_CHUNK) {
        if (!net->chunk_cache) return false;

        pxl8_chunk_msg_header chunk_hdr;
        offset += pxl8_protocol_deserialize_chunk_msg_header(pkt->data + offset, pkt->len - offset, &chunk_hdr);

        const u8* payload = pkt->data + offset;
        usize payload_len = chunk_hdr.payload_size;
        if (payload_len > pkt->len - offset) {
            payload_len = pkt->len - offset;
        }

        pxl8_world_chunk_cache_receive(net->chunk_cache, &chunk_hdr, payload, payload_len);
        return true;
    }

    if (hdr.type == PXL8_MSG_CHUNK_ENTER) {
        pxl8_chunk_enter_msg chunk_msg;
        pxl8_protocol_deserialize_chunk_enter(pkt->data + offset, pkt->len - offset, &chunk_msg);
        net->chunk_id = chunk_msg.chunk_id;
        net->chunk_type = chunk_msg.chunk_type;
        return true;
    }

    if (hdr.type == PXL8_MSG_CHUNK_EXIT) {
        net->chunk_id = 0;
        net->chunk_type = 0;
        return true;
    }

    if (hdr.type != PXL8_MSG_SNAPSHOT) return false;

    pxl8_snapshot_header snap;
    offset += pxl8_protocol_deserialize_snapshot_header(pkt->data + offset, pkt->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(
            pkt->data + offset, pkt->len - offset, &net->entities[i]);
    }

    return true;
}

#endif