#include "pxl8_entity.h"

#include <string.h>

#include "pxl8_mem.h"

#define PXL8_ENTITY_COMPONENT_NAME_MAX 32
#define PXL8_ENTITY_RELATIONSHIP_NAME_MAX 32

typedef struct pxl8_component_type {
    char name[PXL8_ENTITY_COMPONENT_NAME_MAX];
    u32 size;
} pxl8_component_type;

typedef struct pxl8_component_storage {
    u32* sparse;
    void* dense_data;
    pxl8_entity* dense_entities;
    u32 count;
} pxl8_component_storage;

typedef struct pxl8_relationship_type {
    char name[PXL8_ENTITY_RELATIONSHIP_NAME_MAX];
} pxl8_relationship_type;

typedef struct pxl8_relationship_entry {
    pxl8_entity subject;
    pxl8_entity object;
    pxl8_entity_relationship rel;
    u32 next_by_subject;
    u32 next_by_object;
} pxl8_relationship_entry;

struct pxl8_entity_pool {
    u32* generations;
    u32* free_list;
    u32 free_count;
    u32 capacity;
    u32 alive_count;

    pxl8_component_type* component_types;
    pxl8_component_storage* component_storage;
    u32 component_type_count;
    u32 component_type_capacity;

    pxl8_relationship_type* relationship_types;
    u32 relationship_type_count;
    u32 relationship_type_capacity;

    pxl8_relationship_entry* relationships;
    u32* rel_by_subject;
    u32* rel_by_object;
    u32 relationship_count;
    u32 relationship_capacity;
};

pxl8_entity_pool* pxl8_entity_pool_create(u32 capacity) {
    pxl8_entity_pool* pool = pxl8_calloc(1, sizeof(pxl8_entity_pool));
    if (!pool) return NULL;

    pool->capacity = capacity;
    pool->generations = pxl8_calloc(capacity, sizeof(u32));
    pool->free_list = pxl8_malloc(capacity * sizeof(u32));

    if (!pool->generations || !pool->free_list) {
        pxl8_entity_pool_destroy(pool);
        return NULL;
    }

    for (u32 i = 0; i < capacity; i++) {
        pool->free_list[i] = capacity - 1 - i;
    }
    pool->free_count = capacity;

    pool->component_type_capacity = 16;
    pool->component_types = pxl8_calloc(pool->component_type_capacity, sizeof(pxl8_component_type));
    pool->component_storage = pxl8_calloc(pool->component_type_capacity, sizeof(pxl8_component_storage));

    pool->relationship_type_capacity = 16;
    pool->relationship_types = pxl8_calloc(pool->relationship_type_capacity, sizeof(pxl8_relationship_type));

    pool->relationship_capacity = 256;
    pool->relationships = pxl8_malloc(pool->relationship_capacity * sizeof(pxl8_relationship_entry));
    pool->rel_by_subject = pxl8_malloc(capacity * sizeof(u32));
    pool->rel_by_object = pxl8_malloc(capacity * sizeof(u32));

    for (u32 i = 0; i < capacity; i++) {
        pool->rel_by_subject[i] = UINT32_MAX;
        pool->rel_by_object[i] = UINT32_MAX;
    }

    return pool;
}

void pxl8_entity_pool_clear(pxl8_entity_pool* pool) {
    if (!pool) return;

    for (u32 i = 0; i < pool->capacity; i++) {
        pool->generations[i] = 0;
        pool->free_list[i] = pool->capacity - 1 - i;
        pool->rel_by_subject[i] = UINT32_MAX;
        pool->rel_by_object[i] = UINT32_MAX;
    }
    pool->free_count = pool->capacity;
    pool->alive_count = 0;

    for (u32 i = 0; i < pool->component_type_count; i++) {
        pool->component_storage[i].count = 0;
    }

    pool->relationship_count = 0;
}

void pxl8_entity_pool_destroy(pxl8_entity_pool* pool) {
    if (!pool) return;

    for (u32 i = 0; i < pool->component_type_count; i++) {
        pxl8_free(pool->component_storage[i].sparse);
        pxl8_free(pool->component_storage[i].dense_data);
        pxl8_free(pool->component_storage[i].dense_entities);
    }

    pxl8_free(pool->component_types);
    pxl8_free(pool->component_storage);
    pxl8_free(pool->relationship_types);
    pxl8_free(pool->relationships);
    pxl8_free(pool->rel_by_subject);
    pxl8_free(pool->rel_by_object);
    pxl8_free(pool->generations);
    pxl8_free(pool->free_list);
    pxl8_free(pool);
}

pxl8_entity pxl8_entity_spawn(pxl8_entity_pool* pool) {
    if (!pool || pool->free_count == 0) return PXL8_ENTITY_INVALID;

    u32 idx = pool->free_list[--pool->free_count];
    pool->generations[idx]++;
    pool->alive_count++;

    return (pxl8_entity){ .idx = idx, .gen = pool->generations[idx] };
}

void pxl8_entity_despawn(pxl8_entity_pool* pool, pxl8_entity e) {
    if (!pool || !pxl8_entity_alive(pool, e)) return;

    for (u32 i = 0; i < pool->component_type_count; i++) {
        pxl8_entity_component_remove(pool, e, i + 1);
    }

    pool->free_list[pool->free_count++] = e.idx;
    pool->generations[e.idx]++;
    pool->alive_count--;
}

bool pxl8_entity_alive(const pxl8_entity_pool* pool, pxl8_entity e) {
    if (!pool || e.idx >= pool->capacity) return false;
    return pool->generations[e.idx] == e.gen && e.gen != 0;
}

u32 pxl8_entity_count(const pxl8_entity_pool* pool) {
    return pool ? pool->alive_count : 0;
}

pxl8_entity_component pxl8_entity_component_register(pxl8_entity_pool* pool, const char* name, u32 size) {
    if (!pool || !name || size == 0) return PXL8_ENTITY_COMPONENT_INVALID;

    pxl8_entity_component existing = pxl8_entity_component_find(pool, name);
    if (existing != PXL8_ENTITY_COMPONENT_INVALID) return existing;

    if (pool->component_type_count >= pool->component_type_capacity) {
        u32 new_capacity = pool->component_type_capacity * 2;
        pxl8_component_type* new_types = pxl8_realloc(pool->component_types, new_capacity * sizeof(pxl8_component_type));
        pxl8_component_storage* new_storage = pxl8_realloc(pool->component_storage, new_capacity * sizeof(pxl8_component_storage));
        if (!new_types || !new_storage) return PXL8_ENTITY_COMPONENT_INVALID;

        pool->component_types = new_types;
        pool->component_storage = new_storage;
        memset(&pool->component_types[pool->component_type_capacity], 0, (new_capacity - pool->component_type_capacity) * sizeof(pxl8_component_type));
        memset(&pool->component_storage[pool->component_type_capacity], 0, (new_capacity - pool->component_type_capacity) * sizeof(pxl8_component_storage));
        pool->component_type_capacity = new_capacity;
    }

    u32 type_idx = pool->component_type_count++;
    strncpy(pool->component_types[type_idx].name, name, PXL8_ENTITY_COMPONENT_NAME_MAX - 1);
    pool->component_types[type_idx].size = size;

    pxl8_component_storage* storage = &pool->component_storage[type_idx];
    storage->sparse = pxl8_malloc(pool->capacity * sizeof(u32));
    storage->dense_data = pxl8_malloc(pool->capacity * size);
    storage->dense_entities = pxl8_malloc(pool->capacity * sizeof(pxl8_entity));
    storage->count = 0;

    for (u32 i = 0; i < pool->capacity; i++) {
        storage->sparse[i] = UINT32_MAX;
    }

    return type_idx + 1;
}

pxl8_entity_component pxl8_entity_component_find(const pxl8_entity_pool* pool, const char* name) {
    if (!pool || !name) return PXL8_ENTITY_COMPONENT_INVALID;

    for (u32 i = 0; i < pool->component_type_count; i++) {
        if (strcmp(pool->component_types[i].name, name) == 0) {
            return i + 1;
        }
    }
    return PXL8_ENTITY_COMPONENT_INVALID;
}

const char* pxl8_entity_component_name(const pxl8_entity_pool* pool, pxl8_entity_component comp) {
    if (!pool || comp == 0 || comp > pool->component_type_count) return NULL;
    return pool->component_types[comp - 1].name;
}

void* pxl8_entity_component_add(pxl8_entity_pool* pool, pxl8_entity e, pxl8_entity_component comp) {
    if (!pool || !pxl8_entity_alive(pool, e)) return NULL;
    if (comp == 0 || comp > pool->component_type_count) return NULL;

    u32 type_idx = comp - 1;
    pxl8_component_storage* storage = &pool->component_storage[type_idx];
    u32 size = pool->component_types[type_idx].size;

    if (storage->sparse[e.idx] != UINT32_MAX) {
        return (u8*)storage->dense_data + storage->sparse[e.idx] * size;
    }

    u32 dense_idx = storage->count++;
    storage->sparse[e.idx] = dense_idx;
    storage->dense_entities[dense_idx] = e;

    void* data = (u8*)storage->dense_data + dense_idx * size;
    memset(data, 0, size);
    return data;
}

void* pxl8_entity_component_get(const pxl8_entity_pool* pool, pxl8_entity e, pxl8_entity_component comp) {
    if (!pool || !pxl8_entity_alive(pool, e)) return NULL;
    if (comp == 0 || comp > pool->component_type_count) return NULL;

    u32 type_idx = comp - 1;
    const pxl8_component_storage* storage = &pool->component_storage[type_idx];

    if (storage->sparse[e.idx] == UINT32_MAX) return NULL;

    u32 size = pool->component_types[type_idx].size;
    return (u8*)storage->dense_data + storage->sparse[e.idx] * size;
}

void pxl8_entity_component_remove(pxl8_entity_pool* pool, pxl8_entity e, pxl8_entity_component comp) {
    if (!pool || !pxl8_entity_alive(pool, e)) return;
    if (comp == 0 || comp > pool->component_type_count) return;

    u32 type_idx = comp - 1;
    pxl8_component_storage* storage = &pool->component_storage[type_idx];
    u32 size = pool->component_types[type_idx].size;

    u32 dense_idx = storage->sparse[e.idx];
    if (dense_idx == UINT32_MAX) return;

    u32 last_idx = storage->count - 1;
    if (dense_idx != last_idx) {
        pxl8_entity last_entity = storage->dense_entities[last_idx];
        memcpy((u8*)storage->dense_data + dense_idx * size,
               (u8*)storage->dense_data + last_idx * size, size);
        storage->dense_entities[dense_idx] = last_entity;
        storage->sparse[last_entity.idx] = dense_idx;
    }

    storage->sparse[e.idx] = UINT32_MAX;
    storage->count--;
}

bool pxl8_entity_component_has(const pxl8_entity_pool* pool, pxl8_entity e, pxl8_entity_component comp) {
    if (!pool || !pxl8_entity_alive(pool, e)) return false;
    if (comp == 0 || comp > pool->component_type_count) return false;

    return pool->component_storage[comp - 1].sparse[e.idx] != UINT32_MAX;
}

pxl8_entity_relationship pxl8_entity_relationship_register(pxl8_entity_pool* pool, const char* name) {
    if (!pool || !name) return PXL8_ENTITY_RELATIONSHIP_INVALID;

    pxl8_entity_relationship existing = pxl8_entity_relationship_find(pool, name);
    if (existing != PXL8_ENTITY_RELATIONSHIP_INVALID) return existing;

    if (pool->relationship_type_count >= pool->relationship_type_capacity) {
        u32 new_capacity = pool->relationship_type_capacity * 2;
        pxl8_relationship_type* new_types = pxl8_realloc(pool->relationship_types, new_capacity * sizeof(pxl8_relationship_type));
        if (!new_types) return PXL8_ENTITY_RELATIONSHIP_INVALID;

        pool->relationship_types = new_types;
        memset(&pool->relationship_types[pool->relationship_type_capacity], 0, (new_capacity - pool->relationship_type_capacity) * sizeof(pxl8_relationship_type));
        pool->relationship_type_capacity = new_capacity;
    }

    u32 type_idx = pool->relationship_type_count++;
    strncpy(pool->relationship_types[type_idx].name, name, PXL8_ENTITY_RELATIONSHIP_NAME_MAX - 1);

    return type_idx + 1;
}

pxl8_entity_relationship pxl8_entity_relationship_find(const pxl8_entity_pool* pool, const char* name) {
    if (!pool || !name) return PXL8_ENTITY_RELATIONSHIP_INVALID;

    for (u32 i = 0; i < pool->relationship_type_count; i++) {
        if (strcmp(pool->relationship_types[i].name, name) == 0) {
            return i + 1;
        }
    }
    return PXL8_ENTITY_RELATIONSHIP_INVALID;
}

const char* pxl8_entity_relationship_name(const pxl8_entity_pool* pool, pxl8_entity_relationship rel) {
    if (!pool || rel == 0 || rel > pool->relationship_type_count) return NULL;
    return pool->relationship_types[rel - 1].name;
}

void pxl8_entity_relationship_add(pxl8_entity_pool* pool, pxl8_entity subject, pxl8_entity_relationship rel, pxl8_entity object) {
    if (!pool) return;
    if (!pxl8_entity_alive(pool, subject) || !pxl8_entity_alive(pool, object)) return;
    if (rel == 0 || rel > pool->relationship_type_count) return;
    if (pxl8_entity_relationship_has(pool, subject, rel, object)) return;

    if (pool->relationship_count >= pool->relationship_capacity) {
        u32 new_capacity = pool->relationship_capacity * 2;
        pxl8_relationship_entry* new_rels = pxl8_realloc(pool->relationships, new_capacity * sizeof(pxl8_relationship_entry));
        if (!new_rels) return;
        pool->relationships = new_rels;
        pool->relationship_capacity = new_capacity;
    }

    u32 entry_idx = pool->relationship_count++;
    pxl8_relationship_entry* entry = &pool->relationships[entry_idx];
    entry->subject = subject;
    entry->object = object;
    entry->rel = rel;

    entry->next_by_subject = pool->rel_by_subject[subject.idx];
    pool->rel_by_subject[subject.idx] = entry_idx;

    entry->next_by_object = pool->rel_by_object[object.idx];
    pool->rel_by_object[object.idx] = entry_idx;
}

void pxl8_entity_relationship_remove(pxl8_entity_pool* pool, pxl8_entity subject, pxl8_entity_relationship rel, pxl8_entity object) {
    if (!pool) return;
    if (rel == 0 || rel > pool->relationship_type_count) return;

    u32* prev_ptr = &pool->rel_by_subject[subject.idx];
    u32 idx = *prev_ptr;

    while (idx != UINT32_MAX) {
        pxl8_relationship_entry* entry = &pool->relationships[idx];
        if (pxl8_entity_eq(entry->subject, subject) &&
            pxl8_entity_eq(entry->object, object) &&
            entry->rel == rel) {

            *prev_ptr = entry->next_by_subject;

            u32* obj_prev = &pool->rel_by_object[object.idx];
            while (*obj_prev != UINT32_MAX) {
                if (*obj_prev == idx) {
                    *obj_prev = entry->next_by_object;
                    break;
                }
                obj_prev = &pool->relationships[*obj_prev].next_by_object;
            }

            if (idx != pool->relationship_count - 1) {
                u32 last_idx = pool->relationship_count - 1;
                pxl8_relationship_entry* last = &pool->relationships[last_idx];

                u32* last_subj_prev = &pool->rel_by_subject[last->subject.idx];
                while (*last_subj_prev != UINT32_MAX) {
                    if (*last_subj_prev == last_idx) {
                        *last_subj_prev = idx;
                        break;
                    }
                    last_subj_prev = &pool->relationships[*last_subj_prev].next_by_subject;
                }

                u32* last_obj_prev = &pool->rel_by_object[last->object.idx];
                while (*last_obj_prev != UINT32_MAX) {
                    if (*last_obj_prev == last_idx) {
                        *last_obj_prev = idx;
                        break;
                    }
                    last_obj_prev = &pool->relationships[*last_obj_prev].next_by_object;
                }

                *entry = *last;
            }
            pool->relationship_count--;
            return;
        }
        prev_ptr = &entry->next_by_subject;
        idx = *prev_ptr;
    }
}

bool pxl8_entity_relationship_has(const pxl8_entity_pool* pool, pxl8_entity subject, pxl8_entity_relationship rel, pxl8_entity object) {
    if (!pool) return false;
    if (rel == 0 || rel > pool->relationship_type_count) return false;

    u32 idx = pool->rel_by_subject[subject.idx];
    while (idx != UINT32_MAX) {
        const pxl8_relationship_entry* entry = &pool->relationships[idx];
        if (pxl8_entity_eq(entry->subject, subject) &&
            pxl8_entity_eq(entry->object, object) &&
            entry->rel == rel) {
            return true;
        }
        idx = entry->next_by_subject;
    }
    return false;
}

u32 pxl8_entity_relationship_subjects(const pxl8_entity_pool* pool, pxl8_entity object, pxl8_entity_relationship rel, pxl8_entity* out, u32 max) {
    if (!pool || !out || max == 0) return 0;
    if (rel == 0 || rel > pool->relationship_type_count) return 0;

    u32 count = 0;
    u32 idx = pool->rel_by_object[object.idx];
    while (idx != UINT32_MAX && count < max) {
        const pxl8_relationship_entry* entry = &pool->relationships[idx];
        if (pxl8_entity_eq(entry->object, object) && entry->rel == rel) {
            out[count++] = entry->subject;
        }
        idx = entry->next_by_object;
    }
    return count;
}

u32 pxl8_entity_relationship_objects(const pxl8_entity_pool* pool, pxl8_entity subject, pxl8_entity_relationship rel, pxl8_entity* out, u32 max) {
    if (!pool || !out || max == 0) return 0;
    if (rel == 0 || rel > pool->relationship_type_count) return 0;

    u32 count = 0;
    u32 idx = pool->rel_by_subject[subject.idx];
    while (idx != UINT32_MAX && count < max) {
        const pxl8_relationship_entry* entry = &pool->relationships[idx];
        if (pxl8_entity_eq(entry->subject, subject) && entry->rel == rel) {
            out[count++] = entry->object;
        }
        idx = entry->next_by_subject;
    }
    return count;
}

void pxl8_entity_each(pxl8_entity_pool* pool, pxl8_entity_component comp, pxl8_entity_each_fn fn, void* ctx) {
    if (!pool || !fn) return;
    if (comp == 0 || comp > pool->component_type_count) return;

    pxl8_component_storage* storage = &pool->component_storage[comp - 1];
    for (u32 i = 0; i < storage->count; i++) {
        pxl8_entity e = storage->dense_entities[i];
        if (pxl8_entity_alive(pool, e)) {
            fn(pool, e, ctx);
        }
    }
}

void pxl8_entity_each_with(pxl8_entity_pool* pool, const pxl8_entity_component* comps, u32 count, pxl8_entity_each_fn fn, void* ctx) {
    if (!pool || !comps || !fn || count == 0) return;

    u32 smallest_idx = 0;
    u32 smallest_count = UINT32_MAX;

    for (u32 i = 0; i < count; i++) {
        if (comps[i] == 0 || comps[i] > pool->component_type_count) return;
        u32 storage_count = pool->component_storage[comps[i] - 1].count;
        if (storage_count < smallest_count) {
            smallest_count = storage_count;
            smallest_idx = i;
        }
    }

    pxl8_component_storage* storage = &pool->component_storage[comps[smallest_idx] - 1];
    for (u32 i = 0; i < storage->count; i++) {
        pxl8_entity e = storage->dense_entities[i];
        if (!pxl8_entity_alive(pool, e)) continue;

        bool has_all = true;
        for (u32 j = 0; j < count && has_all; j++) {
            if (j != smallest_idx) {
                has_all = pxl8_entity_component_has(pool, e, comps[j]);
            }
        }

        if (has_all) {
            fn(pool, e, ctx);
        }
    }
}