#include "pxl8_platform.h"
#include "demo3d_net.h"

#include <string.h>

#include "pxl8_bytes.h"
#include "pxl8_log.h"
#include "pxl8_mem.h"
#include "demo3d_world.h"
#include "demo3d_chunk_cache.h"

static const demo3d_entity_state* find_entity(const demo3d_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;
}

static bool dispatch_message(demo3d_net* ng, const u8* data, usize len) {
    if (len < sizeof(demo3d_msg_header)) return false;

    demo3d_msg_header hdr;
    usize offset = demo3d_protocol_deserialize_header(data, len, &hdr);

    if (hdr.type == DEMO3D_MSG_CHUNK) {
        if (!ng->chunk_cache) return false;
        demo3d_chunk_msg_header chunk_hdr;
        offset += demo3d_protocol_deserialize_chunk_msg_header(data + offset, len - offset, &chunk_hdr);
        const u8* payload = data + offset;
        usize payload_len = chunk_hdr.payload_size;
        if (payload_len > len - offset) payload_len = len - offset;
        demo3d_chunk_cache_receive(ng->chunk_cache, &chunk_hdr, payload, payload_len);
        return true;
    }

    if (hdr.type == DEMO3D_MSG_CHUNK_ENTER) {
        demo3d_chunk_enter_msg chunk_msg;
        demo3d_protocol_deserialize_chunk_enter(data + offset, len - offset, &chunk_msg);
        ng->chunk_cx = chunk_msg.cx;
        ng->chunk_cz = chunk_msg.cz;
        ng->has_chunk = true;
        pxl8_debug("Received CHUNK_ENTER cx=%d cz=%d", chunk_msg.cx, chunk_msg.cz);
        return true;
    }

    if (hdr.type == DEMO3D_MSG_CHUNK_EXIT) {
        ng->has_chunk = false;
        return true;
    }

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

    demo3d_snapshot_header snap;
    offset += demo3d_protocol_deserialize_snapshot_header(data + offset, len - offset, &snap);
    if (snap.tick <= ng->highest_tick) return false;

    memcpy(ng->prev_entities, ng->entities, sizeof(ng->entities));
    ng->prev_snapshot = ng->snapshot;
    ng->highest_tick = snap.tick;
    ng->snapshot = snap;
    ng->interp_time = 0.0f;

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

    for (u16 i = 0; i < count; i++) {
        offset += demo3d_protocol_deserialize_entity_state(data + offset, len - offset, &ng->entities[i]);
    }

    return true;
}

demo3d_net* demo3d_net_create(const pxl8_net_config* config) {
    demo3d_net* ng = pxl8_calloc(1, sizeof(demo3d_net));
    if (!ng) return NULL;
    ng->transport = pxl8_net_create(config);
    if (!ng->transport) {
        pxl8_free(ng);
        return NULL;
    }
    return ng;
}

void demo3d_net_destroy(demo3d_net* ng) {
    if (!ng) return;
    pxl8_net_destroy(ng->transport);
    pxl8_free(ng);
}

pxl8_result demo3d_net_connect(demo3d_net* ng) {
    if (!ng) return PXL8_ERROR_INVALID_ARGUMENT;
    return pxl8_net_connect(ng->transport);
}

bool demo3d_net_connected(const demo3d_net* ng) {
    return ng && pxl8_net_connected(ng->transport);
}

void demo3d_net_disconnect(demo3d_net* ng) {
    if (!ng) return;
    pxl8_net_disconnect(ng->transport);
}

const demo3d_entity_state* demo3d_net_entities(const demo3d_net* ng) {
    if (!ng) return NULL;
    return ng->entities;
}

const u8* demo3d_net_entity_prev_userdata(const demo3d_net* ng, u64 entity_id) {
    if (!ng) return NULL;
    const demo3d_entity_state* e = find_entity(ng->prev_entities, ng->prev_snapshot.entity_count, entity_id);
    return e ? e->userdata : NULL;
}

const u8* demo3d_net_entity_userdata(const demo3d_net* ng, u64 entity_id) {
    if (!ng) return NULL;
    const demo3d_entity_state* e = find_entity(ng->entities, ng->snapshot.entity_count, entity_id);
    return e ? e->userdata : NULL;
}

const demo3d_input_msg* demo3d_net_input_at(const demo3d_net* ng, u64 tick) {
    if (!ng) return NULL;
    if (tick < ng->input_oldest_tick) return NULL;
    u64 age = tick - ng->input_oldest_tick;
    if (age >= DEMO3D_NET_INPUT_HISTORY_SIZE) return NULL;
    u64 idx = tick % DEMO3D_NET_INPUT_HISTORY_SIZE;
    const demo3d_input_msg* msg = &ng->input_history[idx];
    if (msg->tick != tick) return NULL;
    return msg;
}

u64 demo3d_net_input_oldest_tick(const demo3d_net* ng) {
    if (!ng) return 0;
    return ng->input_oldest_tick;
}

void demo3d_net_input_push(demo3d_net* ng, const demo3d_input_msg* input) {
    if (!ng || !input) return;
    u64 idx = input->tick % DEMO3D_NET_INPUT_HISTORY_SIZE;
    ng->input_history[idx] = *input;
    ng->input_head = input->tick;
    if (input->tick >= DEMO3D_NET_INPUT_HISTORY_SIZE) {
        ng->input_oldest_tick = input->tick - DEMO3D_NET_INPUT_HISTORY_SIZE + 1;
    }
}

f32 demo3d_net_lerp_alpha(const demo3d_net* ng) {
    if (!ng) return 1.0f;
    return ng->interp_time / (1.0f / DEMO3D_NET_TICK_RATE);
}

bool demo3d_net_needs_correction(const demo3d_net* ng) {
    if (!ng) return false;
    if (ng->snapshot.tick == 0) return false;
    if (ng->predicted_tick == 0) return false;

    u64 player_id = ng->snapshot.player_id;
    const demo3d_entity_state* server = find_entity(ng->entities, ng->snapshot.entity_count, player_id);
    if (!server) return false;

    return memcmp(server->userdata, ng->predicted_state, DEMO3D_NET_USERDATA_SIZE) != 0;
}

u64 demo3d_net_player_id(const demo3d_net* ng) {
    if (!ng) return 0;
    return ng->snapshot.player_id;
}

bool demo3d_net_poll(demo3d_net* ng) {
    if (!ng || !pxl8_net_connected(ng->transport)) return false;

    u8 recv_buf[4096];
    usize len = pxl8_net_recv(ng->transport, recv_buf, sizeof(recv_buf));
    u64 prev_tick = ng->highest_tick;
    if (!dispatch_message(ng, recv_buf, len)) return false;

    if (ng->highest_tick > prev_tick && ng->world) {
        demo3d_world_reconcile(ng->world, ng, 1.0f / 30.0f);
    }

    return true;
}

u8* demo3d_net_predicted_state(demo3d_net* ng) {
    if (!ng) return NULL;
    return ng->predicted_state;
}

void demo3d_net_predicted_tick_set(demo3d_net* ng, u64 tick) {
    if (!ng) return;
    ng->predicted_tick = tick;
}

pxl8_result demo3d_net_send_command(demo3d_net* ng, const demo3d_command_msg* cmd) {
    if (!ng || !cmd) return PXL8_ERROR_INVALID_ARGUMENT;
    if (!pxl8_net_connected(ng->transport)) return PXL8_ERROR_INVALID_ARGUMENT;

    u8 buf[512];
    demo3d_msg_header hdr = {.type = DEMO3D_MSG_COMMAND, .version = DEMO3D_PROTOCOL_VERSION};
    usize offset = demo3d_protocol_serialize_header(&hdr, buf, sizeof(buf));
    offset += demo3d_protocol_serialize_command(cmd, buf + offset, sizeof(buf) - offset);

    return pxl8_net_send(ng->transport, buf, offset);
}

pxl8_result demo3d_net_send_input(demo3d_net* ng, const demo3d_input_msg* input) {
    if (!ng || !input) return PXL8_ERROR_INVALID_ARGUMENT;
    if (!pxl8_net_connected(ng->transport)) return PXL8_ERROR_INVALID_ARGUMENT;

    demo3d_net_input_push(ng, input);

    u8 buf[512];
    demo3d_msg_header hdr = {.type = DEMO3D_MSG_INPUT, .version = DEMO3D_PROTOCOL_VERSION};
    usize offset = demo3d_protocol_serialize_header(&hdr, buf, sizeof(buf));
    offset += demo3d_protocol_serialize_input(input, buf + offset, sizeof(buf) - offset);

    return pxl8_net_send(ng->transport, buf, offset);
}

const demo3d_snapshot_header* demo3d_net_snapshot(const demo3d_net* ng) {
    if (!ng) return NULL;
    return &ng->snapshot;
}

u64 demo3d_net_tick(const demo3d_net* ng) {
    if (!ng) return 0;
    return ng->snapshot.tick;
}

void demo3d_net_update(demo3d_net* ng, f32 dt) {
    if (!ng) return;
    ng->interp_time += dt;
}

void demo3d_net_set_chunk_cache(demo3d_net* ng, demo3d_chunk_cache* cache) {
    if (!ng) return;
    ng->chunk_cache = cache;
}

demo3d_chunk_cache* demo3d_net_chunk_cache(demo3d_net* ng) {
    if (!ng) return NULL;
    return ng->chunk_cache;
}

void demo3d_net_set_world(demo3d_net* ng, demo3d_world* world) {
    if (!ng) return;
    ng->world = world;
}

i32 demo3d_net_chunk_cx(const demo3d_net* ng) {
    if (!ng) return 0;
    return ng->chunk_cx;
}

i32 demo3d_net_chunk_cz(const demo3d_net* ng) {
    if (!ng) return 0;
    return ng->chunk_cz;
}

bool demo3d_net_has_chunk(const demo3d_net* ng) {
    if (!ng) return false;
    return ng->has_chunk;
}

pxl8_result demo3d_net_spawn(demo3d_net* ng, f32 x, f32 y, f32 z, f32 yaw, f32 pitch) {
    if (!ng) return PXL8_ERROR_INVALID_ARGUMENT;

    demo3d_command_msg cmd = {0};
    cmd.cmd_type = DEMO3D_CMD_SPAWN_ENTITY;

    u8* p = cmd.payload;
    memcpy(p, &x, 4); p += 4;
    memcpy(p, &y, 4); p += 4;
    memcpy(p, &z, 4); p += 4;
    memcpy(p, &yaw, 4); p += 4;
    memcpy(p, &pitch, 4);
    cmd.payload_size = 20;

    return demo3d_net_send_command(ng, &cmd);
}

#ifdef PXL8_ASYNC_THREADS
void demo3d_net_start_thread(demo3d_net* ng) {
    if (!ng) return;
    pxl8_net_start_thread(ng->transport);
}

void demo3d_net_stop_thread(demo3d_net* ng) {
    if (!ng) return;
    pxl8_net_stop_thread(ng->transport);
}

bool demo3d_net_dispatch_packet(demo3d_net* ng, const pxl8_packet* pkt) {
    if (!ng || !pkt) return false;
    return dispatch_message(ng, pkt->data, pkt->len);
}
#endif