#include "pxl8_world_chunk_cache.h"

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

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

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->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);
    }
    memset(e, 0, sizeof(*e));
    return e;
}

static void assembly_reset(pxl8_world_chunk_assembly* a) {
    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->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 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, 48, &wire_hdr);
    s.offset = 48;

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

    if (wire_hdr.num_heightfield > 0) {
        bsp->heightfield = pxl8_calloc(wire_hdr.num_heightfield, sizeof(f32));
        bsp->num_heightfield = wire_hdr.num_heightfield;
        for (u32 i = 0; i < wire_hdr.num_heightfield; i++) {
            u32 raw = pxl8_read_u32_be(&s);
            memcpy(&bsp->heightfield[i], &raw, sizeof(f32));
        }
        bsp->heightfield_w = pxl8_read_u16_be(&s);
        bsp->heightfield_h = pxl8_read_u16_be(&s);
        u32 ox_raw = pxl8_read_u32_be(&s);
        memcpy(&bsp->heightfield_ox, &ox_raw, sizeof(f32));
        u32 oz_raw = pxl8_read_u32_be(&s);
        memcpy(&bsp->heightfield_oz, &oz_raw, sizeof(f32));
        u32 cs_raw = pxl8_read_u32_be(&s);
        memcpy(&bsp->heightfield_cell_size, &cs_raw, sizeof(f32));
    }

    u32 raw;
    raw = pxl8_read_u32_be(&s);
    memcpy(&bsp->bounds_min_x, &raw, sizeof(f32));
    raw = pxl8_read_u32_be(&s);
    memcpy(&bsp->bounds_min_z, &raw, sizeof(f32));
    raw = pxl8_read_u32_be(&s);
    memcpy(&bsp->bounds_max_x, &raw, sizeof(f32));
    raw = pxl8_read_u32_be(&s);
    memcpy(&bsp->bounds_max_z, &raw, sizeof(f32));

    return bsp;
}

static pxl8_result assemble_bsp(pxl8_world_chunk_cache* cache, pxl8_world_chunk_assembly* a) {
    pxl8_bsp* bsp = assembly_to_bsp(a);
    if (!bsp) {
        assembly_reset(a);
        return PXL8_ERROR_INVALID_ARGUMENT;
    }

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

    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 ||
                        a->id != hdr->id ||
                        a->version != hdr->version;

    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;
        return assemble_bsp(cache, a);
    }

    return PXL8_OK;
}

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

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