#include "pxl8_world_chunk_cache.h"

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

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

#define PXL8_VXL_CHUNK_VOLUME (PXL8_VXL_CHUNK_SIZE * PXL8_VXL_CHUNK_SIZE * PXL8_VXL_CHUNK_SIZE)

static pxl8_world_chunk_cache_entry* find_entry_vxl(pxl8_world_chunk_cache* cache, i32 cx, i32 cy, i32 cz) {
    for (u32 i = 0; i < cache->entry_count; i++) {
        pxl8_world_chunk_cache_entry* e = &cache->entries[i];
        if (e->valid && e->chunk && e->chunk->type == PXL8_WORLD_CHUNK_VXL &&
            e->chunk->cx == cx && e->chunk->cy == cy && e->chunk->cz == cz) {
            return e;
        }
    }
    return NULL;
}

static pxl8_world_chunk_cache_entry* find_entry_bsp(pxl8_world_chunk_cache* cache, u32 id) {
    for (u32 i = 0; i < cache->entry_count; i++) {
        pxl8_world_chunk_cache_entry* e = &cache->entries[i];
        if (e->valid && e->chunk && e->chunk->type == PXL8_WORLD_CHUNK_BSP &&
            e->chunk->id == id) {
            return e;
        }
    }
    return NULL;
}

static pxl8_world_chunk_cache_entry* alloc_entry(pxl8_world_chunk_cache* cache) {
    if (cache->entry_count < PXL8_WORLD_CHUNK_CACHE_SIZE) {
        pxl8_world_chunk_cache_entry* e = &cache->entries[cache->entry_count++];
        memset(e, 0, sizeof(*e));
        return e;
    }

    for (u32 i = 0; i < PXL8_WORLD_CHUNK_CACHE_SIZE; i++) {
        if (!cache->entries[i].valid) {
            pxl8_world_chunk_cache_entry* e = &cache->entries[i];
            memset(e, 0, sizeof(*e));
            return e;
        }
    }

    u64 oldest = cache->entries[0].last_used;
    u32 slot = 0;
    for (u32 i = 1; i < PXL8_WORLD_CHUNK_CACHE_SIZE; i++) {
        if (cache->entries[i].last_used < oldest) {
            oldest = cache->entries[i].last_used;
            slot = i;
        }
    }

    pxl8_world_chunk_cache_entry* e = &cache->entries[slot];
    if (e->chunk) pxl8_world_chunk_destroy(e->chunk);
    if (e->mesh) pxl8_mesh_destroy(e->mesh);
    memset(e, 0, sizeof(*e));
    return e;
}

static void assembly_reset(pxl8_world_chunk_assembly* a) {
    a->type = PXL8_WORLD_CHUNK_VXL;
    a->id = 0;
    a->cx = 0;
    a->cy = 0;
    a->cz = 0;
    a->version = 0;
    a->fragment_count = 0;
    a->fragments_received = 0;
    a->data_size = 0;
    a->active = false;
    a->complete = false;
}

static void assembly_init(pxl8_world_chunk_assembly* a, const pxl8_chunk_msg_header* hdr) {
    assembly_reset(a);
    a->type = hdr->chunk_type == PXL8_CHUNK_TYPE_BSP ? PXL8_WORLD_CHUNK_BSP : PXL8_WORLD_CHUNK_VXL;
    a->id = hdr->id;
    a->cx = hdr->cx;
    a->cy = hdr->cy;
    a->cz = hdr->cz;
    a->version = hdr->version;
    a->fragment_count = hdr->fragment_count;
    a->active = true;

    u32 needed = PXL8_CHUNK_MAX_PAYLOAD * hdr->fragment_count;
    if (a->data_capacity < needed) {
        a->data_capacity = needed;
        a->data = pxl8_realloc(a->data, a->data_capacity);
    }
}

static bool rle_decode_voxel(const u8* src, usize src_len, pxl8_vxl_chunk* chunk) {
    u8* linear = pxl8_malloc(PXL8_VXL_CHUNK_VOLUME);
    if (!linear) return false;

    usize src_pos = 0;
    usize dst_pos = 0;

    while (src_pos + 1 < src_len && dst_pos < PXL8_VXL_CHUNK_VOLUME) {
        u8 block = src[src_pos++];
        u8 run_minus_one = src[src_pos++];
        usize run = (usize)run_minus_one + 1;

        for (usize i = 0; i < run && dst_pos < PXL8_VXL_CHUNK_VOLUME; i++) {
            linear[dst_pos++] = block;
        }
    }

    if (dst_pos != PXL8_VXL_CHUNK_VOLUME) {
        pxl8_free(linear);
        return false;
    }

    pxl8_vxl_chunk_from_linear(chunk, linear);
    pxl8_free(linear);
    return true;
}

static pxl8_result deserialize_vertex(pxl8_stream* s, pxl8_bsp_vertex* v) {
    v->position.x = pxl8_read_f32_be(s);
    v->position.y = pxl8_read_f32_be(s);
    v->position.z = pxl8_read_f32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_edge(pxl8_stream* s, pxl8_bsp_edge* e) {
    e->vertex[0] = pxl8_read_u16_be(s);
    e->vertex[1] = pxl8_read_u16_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_plane(pxl8_stream* s, pxl8_bsp_plane* p) {
    p->normal.x = pxl8_read_f32_be(s);
    p->normal.y = pxl8_read_f32_be(s);
    p->normal.z = pxl8_read_f32_be(s);
    p->dist = pxl8_read_f32_be(s);
    p->type = (i32)pxl8_read_u32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_face(pxl8_stream* s, pxl8_bsp_face* f) {
    f->first_edge = pxl8_read_u32_be(s);
    f->lightmap_offset = pxl8_read_u32_be(s);
    f->num_edges = pxl8_read_u16_be(s);
    f->plane_id = pxl8_read_u16_be(s);
    f->side = pxl8_read_u16_be(s);
    pxl8_read_bytes(s, f->styles, 4);
    f->material_id = pxl8_read_u16_be(s);
    f->aabb_min.x = pxl8_read_f32_be(s);
    f->aabb_min.y = pxl8_read_f32_be(s);
    f->aabb_min.z = pxl8_read_f32_be(s);
    f->aabb_max.x = pxl8_read_f32_be(s);
    f->aabb_max.y = pxl8_read_f32_be(s);
    f->aabb_max.z = pxl8_read_f32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_node(pxl8_stream* s, pxl8_bsp_node* n) {
    n->children[0] = (i32)pxl8_read_u32_be(s);
    n->children[1] = (i32)pxl8_read_u32_be(s);
    n->first_face = pxl8_read_u16_be(s);
    n->maxs[0] = (i16)pxl8_read_u16_be(s);
    n->maxs[1] = (i16)pxl8_read_u16_be(s);
    n->maxs[2] = (i16)pxl8_read_u16_be(s);
    n->mins[0] = (i16)pxl8_read_u16_be(s);
    n->mins[1] = (i16)pxl8_read_u16_be(s);
    n->mins[2] = (i16)pxl8_read_u16_be(s);
    n->num_faces = pxl8_read_u16_be(s);
    n->plane_id = pxl8_read_u32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_leaf(pxl8_stream* s, pxl8_bsp_leaf* l) {
    pxl8_read_bytes(s, l->ambient_level, 4);
    l->contents = (i32)pxl8_read_u32_be(s);
    l->first_marksurface = pxl8_read_u16_be(s);
    l->maxs[0] = (i16)pxl8_read_u16_be(s);
    l->maxs[1] = (i16)pxl8_read_u16_be(s);
    l->maxs[2] = (i16)pxl8_read_u16_be(s);
    l->mins[0] = (i16)pxl8_read_u16_be(s);
    l->mins[1] = (i16)pxl8_read_u16_be(s);
    l->mins[2] = (i16)pxl8_read_u16_be(s);
    l->num_marksurfaces = pxl8_read_u16_be(s);
    l->visofs = (i32)pxl8_read_u32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_portal(pxl8_stream* s, pxl8_bsp_portal* p) {
    p->x0 = pxl8_read_f32_be(s);
    p->z0 = pxl8_read_f32_be(s);
    p->x1 = pxl8_read_f32_be(s);
    p->z1 = pxl8_read_f32_be(s);
    p->target_leaf = pxl8_read_u32_be(s);
    return PXL8_OK;
}

static pxl8_result deserialize_cell_portals(pxl8_stream* s, pxl8_bsp_cell_portals* cp) {
    cp->num_portals = pxl8_read_u8(s);
    pxl8_read_u8(s);
    pxl8_read_u8(s);
    pxl8_read_u8(s);
    for (int i = 0; i < 4; i++) {
        deserialize_portal(s, &cp->portals[i]);
    }
    return PXL8_OK;
}

static pxl8_bsp* assembly_to_bsp(pxl8_world_chunk_assembly* a) {
    if (!a->complete || a->data_size < 48) {
        return NULL;
    }

    pxl8_bsp* bsp = pxl8_calloc(1, sizeof(pxl8_bsp));
    if (!bsp) return NULL;

    pxl8_stream s = pxl8_stream_create(a->data, (u32)a->data_size);

    pxl8_bsp_wire_header wire_hdr;
    pxl8_protocol_deserialize_bsp_wire_header(a->data, 44, &wire_hdr);
    s.offset = 44;

    pxl8_debug("[CLIENT] Wire header: verts=%u edges=%u faces=%u planes=%u nodes=%u leafs=%u surfedges=%u visdata=%u",
               wire_hdr.num_vertices, wire_hdr.num_edges, wire_hdr.num_faces,
               wire_hdr.num_planes, wire_hdr.num_nodes, wire_hdr.num_leafs,
               wire_hdr.num_surfedges, wire_hdr.visdata_size);

    if (wire_hdr.num_vertices > 0) {
        bsp->vertices = pxl8_calloc(wire_hdr.num_vertices, sizeof(pxl8_bsp_vertex));
        bsp->num_vertices = wire_hdr.num_vertices;
        for (u32 i = 0; i < wire_hdr.num_vertices; i++) {
            deserialize_vertex(&s, &bsp->vertices[i]);
        }
    }

    if (wire_hdr.num_edges > 0) {
        bsp->edges = pxl8_calloc(wire_hdr.num_edges, sizeof(pxl8_bsp_edge));
        bsp->num_edges = wire_hdr.num_edges;
        for (u32 i = 0; i < wire_hdr.num_edges; i++) {
            deserialize_edge(&s, &bsp->edges[i]);
        }
    }

    if (wire_hdr.num_surfedges > 0) {
        bsp->surfedges = pxl8_calloc(wire_hdr.num_surfedges, sizeof(i32));
        bsp->num_surfedges = wire_hdr.num_surfedges;
        for (u32 i = 0; i < wire_hdr.num_surfedges; i++) {
            bsp->surfedges[i] = (i32)pxl8_read_u32_be(&s);
        }
    }

    if (wire_hdr.num_planes > 0) {
        bsp->planes = pxl8_calloc(wire_hdr.num_planes, sizeof(pxl8_bsp_plane));
        bsp->num_planes = wire_hdr.num_planes;
        for (u32 i = 0; i < wire_hdr.num_planes; i++) {
            deserialize_plane(&s, &bsp->planes[i]);
        }
    }

    if (wire_hdr.num_faces > 0) {
        bsp->faces = pxl8_calloc(wire_hdr.num_faces, sizeof(pxl8_bsp_face));
        bsp->num_faces = wire_hdr.num_faces;
        for (u32 i = 0; i < wire_hdr.num_faces; i++) {
            deserialize_face(&s, &bsp->faces[i]);
        }
    }

    if (wire_hdr.num_nodes > 0) {
        bsp->nodes = pxl8_calloc(wire_hdr.num_nodes, sizeof(pxl8_bsp_node));
        bsp->num_nodes = wire_hdr.num_nodes;
        for (u32 i = 0; i < wire_hdr.num_nodes; i++) {
            deserialize_node(&s, &bsp->nodes[i]);
        }
    }

    if (wire_hdr.num_leafs > 0) {
        bsp->leafs = pxl8_calloc(wire_hdr.num_leafs, sizeof(pxl8_bsp_leaf));
        bsp->num_leafs = wire_hdr.num_leafs;
        for (u32 i = 0; i < wire_hdr.num_leafs; i++) {
            deserialize_leaf(&s, &bsp->leafs[i]);
        }
    }

    if (wire_hdr.num_marksurfaces > 0) {
        bsp->marksurfaces = pxl8_calloc(wire_hdr.num_marksurfaces, sizeof(u16));
        bsp->num_marksurfaces = wire_hdr.num_marksurfaces;
        for (u32 i = 0; i < wire_hdr.num_marksurfaces; i++) {
            bsp->marksurfaces[i] = pxl8_read_u16_be(&s);
        }
    }

    if (wire_hdr.num_cell_portals > 0) {
        bsp->cell_portals = pxl8_calloc(wire_hdr.num_cell_portals, sizeof(pxl8_bsp_cell_portals));
        bsp->num_cell_portals = wire_hdr.num_cell_portals;
        for (u32 i = 0; i < wire_hdr.num_cell_portals; i++) {
            deserialize_cell_portals(&s, &bsp->cell_portals[i]);
        }
    }

    if (wire_hdr.visdata_size > 0) {
        bsp->visdata = pxl8_malloc(wire_hdr.visdata_size);
        bsp->visdata_size = wire_hdr.visdata_size;
        pxl8_read_bytes(&s, bsp->visdata, wire_hdr.visdata_size);
    }

    if (wire_hdr.num_vertex_lights > 0) {
        bsp->vertex_lights = pxl8_calloc(wire_hdr.num_vertex_lights, sizeof(u32));
        bsp->num_vertex_lights = wire_hdr.num_vertex_lights;
        for (u32 i = 0; i < wire_hdr.num_vertex_lights; i++) {
            bsp->vertex_lights[i] = pxl8_read_u32_be(&s);
        }
    }

    pxl8_debug("Deserialized BSP: %u verts, %u faces, %u nodes, %u leafs",
               bsp->num_vertices, bsp->num_faces, bsp->num_nodes, bsp->num_leafs);

    return bsp;
}

static pxl8_result assemble_vxl(pxl8_world_chunk_cache* cache, pxl8_world_chunk_assembly* a) {
    pxl8_world_chunk_cache_entry* entry = find_entry_vxl(cache, a->cx, a->cy, a->cz);
    if (!entry) {
        entry = alloc_entry(cache);
        entry->chunk = pxl8_world_chunk_create_vxl(a->cx, a->cy, a->cz);
        entry->valid = true;
    }

    entry->chunk->version = a->version;
    entry->mesh_dirty = true;
    entry->last_used = cache->frame_counter;

    if (entry->mesh) {
        pxl8_mesh_destroy(entry->mesh);
        entry->mesh = NULL;
    }

    pxl8_vxl_block_clear(entry->chunk->voxels);

    if (!rle_decode_voxel(a->data, a->data_size, entry->chunk->voxels)) {
        pxl8_error("[CLIENT] RLE decode failed for chunk (%d,%d,%d)", a->cx, a->cy, a->cz);
        return PXL8_ERROR_INVALID_ARGUMENT;
    }

    assembly_reset(a);
    return PXL8_OK;
}

static pxl8_result assemble_bsp(pxl8_world_chunk_cache* cache, pxl8_world_chunk_assembly* a) {
    pxl8_debug("[CLIENT] assemble_bsp: id=%u data_size=%zu", a->id, a->data_size);
    pxl8_bsp* bsp = assembly_to_bsp(a);
    if (!bsp) {
        pxl8_debug("[CLIENT] assemble_bsp: assembly_to_bsp returned NULL!");
        assembly_reset(a);
        return PXL8_ERROR_INVALID_ARGUMENT;
    }
    pxl8_debug("[CLIENT] assemble_bsp: BSP created with %u verts %u faces", bsp->num_vertices, bsp->num_faces);

    pxl8_world_chunk_cache_entry* entry = find_entry_bsp(cache, a->id);
    if (entry) {
        if (entry->chunk && entry->chunk->bsp) {
            pxl8_bsp_destroy(entry->chunk->bsp);
        }
    } else {
        entry = alloc_entry(cache);
        entry->chunk = pxl8_world_chunk_create_bsp(a->id);
        entry->valid = true;
    }

    entry->chunk->bsp = bsp;
    entry->chunk->version = a->version;
    entry->last_used = cache->frame_counter;

    assembly_reset(a);
    return PXL8_OK;
}

pxl8_world_chunk_cache* pxl8_world_chunk_cache_create(void) {
    pxl8_world_chunk_cache* cache = pxl8_calloc(1, sizeof(pxl8_world_chunk_cache));
    if (!cache) return NULL;
    assembly_reset(&cache->assembly);
    return cache;
}

void pxl8_world_chunk_cache_destroy(pxl8_world_chunk_cache* cache) {
    if (!cache) return;

    for (u32 i = 0; i < cache->entry_count; i++) {
        pxl8_world_chunk_cache_entry* e = &cache->entries[i];
        if (e->chunk) pxl8_world_chunk_destroy(e->chunk);
        if (e->mesh) pxl8_mesh_destroy(e->mesh);
    }

    pxl8_free(cache->assembly.data);
    pxl8_free(cache);
}

pxl8_result pxl8_world_chunk_cache_receive(pxl8_world_chunk_cache* cache,
                                            const pxl8_chunk_msg_header* hdr,
                                            const u8* payload, usize len) {
    if (!cache || !hdr || !payload) return PXL8_ERROR_INVALID_ARGUMENT;

    pxl8_world_chunk_assembly* a = &cache->assembly;

    bool new_assembly = !a->active ||
                        (hdr->chunk_type == PXL8_CHUNK_TYPE_BSP && a->id != hdr->id) ||
                        (hdr->chunk_type == PXL8_CHUNK_TYPE_VXL &&
                         (a->cx != hdr->cx || a->cy != hdr->cy || a->cz != hdr->cz)) ||
                        a->version != hdr->version ||
                        hdr->fragment_idx == 0;

    if (new_assembly) {
        assembly_init(a, hdr);
    }

    if (hdr->fragment_idx >= PXL8_WORLD_CHUNK_MAX_FRAGMENTS) {
        return PXL8_ERROR_INVALID_ARGUMENT;
    }

    u32 offset = (u32)hdr->fragment_idx * PXL8_CHUNK_MAX_PAYLOAD;
    u32 required = offset + (u32)len;

    if (required > a->data_capacity) {
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    memcpy(a->data + offset, payload, len);

    if (required > a->data_size) {
        a->data_size = required;
    }

    a->fragments_received++;

    if (hdr->flags & PXL8_CHUNK_FLAG_FINAL) {
        a->complete = true;
        if (a->type == PXL8_WORLD_CHUNK_BSP) {
            return assemble_bsp(cache, a);
        } else {
            return assemble_vxl(cache, a);
        }
    }

    return PXL8_OK;
}

pxl8_world_chunk* pxl8_world_chunk_cache_get_vxl(pxl8_world_chunk_cache* cache, i32 cx, i32 cy, i32 cz) {
    if (!cache) return NULL;
    pxl8_world_chunk_cache_entry* e = find_entry_vxl(cache, cx, cy, cz);
    if (e) {
        e->last_used = cache->frame_counter;
        return e->chunk;
    }
    return NULL;
}

pxl8_world_chunk* pxl8_world_chunk_cache_get_bsp(pxl8_world_chunk_cache* cache, u32 id) {
    if (!cache) return NULL;
    pxl8_world_chunk_cache_entry* e = find_entry_bsp(cache, id);
    if (e) {
        e->last_used = cache->frame_counter;
        return e->chunk;
    }
    return NULL;
}

pxl8_mesh* pxl8_world_chunk_cache_get_mesh(pxl8_world_chunk_cache* cache,
                                            i32 cx, i32 cy, i32 cz,
                                            const pxl8_vxl_block_registry* registry,
                                            const pxl8_vxl_mesh_config* config) {
    if (!cache || !registry) return NULL;

    pxl8_world_chunk_cache_entry* entry = find_entry_vxl(cache, cx, cy, cz);
    if (!entry || !entry->chunk || !entry->chunk->voxels) return NULL;

    pxl8_world_chunk* nx = pxl8_world_chunk_cache_get_vxl(cache, cx - 1, cy, cz);
    pxl8_world_chunk* px = pxl8_world_chunk_cache_get_vxl(cache, cx + 1, cy, cz);
    pxl8_world_chunk* ny = pxl8_world_chunk_cache_get_vxl(cache, cx, cy - 1, cz);
    pxl8_world_chunk* py = pxl8_world_chunk_cache_get_vxl(cache, cx, cy + 1, cz);
    pxl8_world_chunk* nz = pxl8_world_chunk_cache_get_vxl(cache, cx, cy, cz - 1);
    pxl8_world_chunk* pz = pxl8_world_chunk_cache_get_vxl(cache, cx, cy, cz + 1);

    bool all_neighbors = nx && px && ny && py && nz && pz;

    if (entry->mesh && !entry->mesh_dirty) {
        if (entry->has_all_neighbors == all_neighbors) {
            return entry->mesh;
        }
    }

    if (entry->mesh) {
        pxl8_mesh_destroy(entry->mesh);
        entry->mesh = NULL;
    }

    const pxl8_vxl_chunk* neighbors[6] = {
        nx ? nx->voxels : NULL,
        px ? px->voxels : NULL,
        ny ? ny->voxels : NULL,
        py ? py->voxels : NULL,
        nz ? nz->voxels : NULL,
        pz ? pz->voxels : NULL
    };

    pxl8_vxl_mesh_config local_config = config ? *config : PXL8_VXL_MESH_CONFIG_DEFAULT;

    entry->mesh = pxl8_vxl_build_mesh(entry->chunk->voxels, neighbors, registry, &local_config);
    entry->mesh_dirty = false;
    entry->has_all_neighbors = all_neighbors;

    return entry->mesh;
}

void pxl8_world_chunk_cache_tick(pxl8_world_chunk_cache* cache) {
    if (!cache) return;
    cache->frame_counter++;
}

void pxl8_world_chunk_cache_evict_distant(pxl8_world_chunk_cache* cache,
                                           i32 cx, i32 cy, i32 cz, i32 radius) {
    if (!cache) return;

    for (u32 i = 0; i < cache->entry_count; i++) {
        pxl8_world_chunk_cache_entry* e = &cache->entries[i];
        if (!e->valid || !e->chunk || e->chunk->type != PXL8_WORLD_CHUNK_VXL) continue;

        i32 dx = e->chunk->cx - cx;
        i32 dy = e->chunk->cy - cy;
        i32 dz = e->chunk->cz - cz;

        if (abs(dx) > radius || abs(dy) > radius || abs(dz) > radius) {
            pxl8_world_chunk_destroy(e->chunk);
            if (e->mesh) pxl8_mesh_destroy(e->mesh);
            e->chunk = NULL;
            e->mesh = NULL;
            e->valid = false;
        }
    }
}

void pxl8_world_chunk_cache_invalidate_meshes(pxl8_world_chunk_cache* cache) {
    if (!cache) return;

    for (u32 i = 0; i < cache->entry_count; i++) {
        cache->entries[i].mesh_dirty = true;
    }
}