/*
 * Usage:
 *
 * Define BECS_MAX_ENTITIES and BECS_MAX_COMPONENTS. Otherwise they will have
 * the following defaults:
 *
 * #define BECS_MAX_ENTITIES 4096
 * #define BECS_MAX_COMPONENTS 64
 *
 * Then define BECS_IMPLEMENTATION in exactly _one_ file before including this
 * file as follows:
 *
 * #define BECS_IMPLEMENTATION
 * #include "becs.h"
 */

#ifndef BECS_H
#define BECS_H

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdbool.h>

#ifndef BECS_NO_ASSERT
#include <assert.h>
#else
#define assert(expression)
#endif /* BECS_SLOW */

typedef struct BECS_Arena BECS_Arena;
typedef struct BECS_ComponentPool BECS_ComponentPool;
typedef struct BECS_ECS BECS_ECS;

typedef uint32_t BECS_Entity;
typedef uint32_t BECS_ComponentID;

size_t BECS_GetMinMemoryArenaSize(size_t *componentSizes, size_t arrayLength);
BECS_ECS BECS_InitECS(void *memory, size_t size);
void BECS_ResetECS(BECS_ECS *ecs);

BECS_Entity BECS_EntityCreate(BECS_ECS *ecs);
void BECS_EntityDestroy(BECS_ECS *ecs, BECS_Entity entity);
bool BECS_IsEntityAlive(BECS_ECS *ecs, BECS_Entity entity);

BECS_ComponentID BECS_ComponentRegister(BECS_ECS *ecs, size_t componentSize, size_t maximumComponentCount);
void *BECS_ComponentAdd(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId);
void BECS_ComponentRemove(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId);
void *BECS_ComponentGet(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId);
bool BECS_ComponentHas(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId);

#ifdef BECS_IMPLEMENTATION

#ifndef BECS_MAX_ENTITIES
#define BECS_MAX_ENTITIES 4096
#endif /* BECS_MAX_ENTITIES */

#ifndef BECS_MAX_COMPONENTS
#define BECS_MAX_COMPONENTS 64
#endif /* BECS_MAX_COMPONENTS */

#define BECS_NULL_ENTITY BECS_MAX_ENTITIES + 1
#define BECS_INVALID_COMPONENT_ID UINT32_MAX

typedef uint64_t BECS_Signature;

struct BECS_Arena {
    uintptr_t nextAllocation;
    size_t capacity;
    void *memory;
};

struct BECS_ComponentPool {
    size_t unitSize;
    size_t capacity;
    size_t count;
    uint32_t sparseArray[BECS_MAX_ENTITIES];
    void *denseArray;
    BECS_Entity *denseEntityMapping;
};

struct BECS_ECS {
    BECS_Arena arena;

    /* Entity Management */
    BECS_Entity nextEntity;
    uint32_t entityCount;
    BECS_Signature entitySignatures[BECS_MAX_ENTITIES];
    bool aliveEntities[BECS_MAX_ENTITIES];

    BECS_Entity freeEntities[BECS_MAX_ENTITIES];
    uint32_t freeEntityCount;

    /* Component Management */
    uint32_t componentTypeCount;
    BECS_ComponentPool componentPools[BECS_MAX_COMPONENTS];
};

/* Signature Helpers */
static BECS_Signature BECS_SignatureSet(BECS_Signature sig, BECS_ComponentID componentId)
{
    return sig | (1ULL << componentId);
}

static bool BECS_SignatureHas(BECS_Signature sig, BECS_ComponentID component_id)
{
    return (sig & (1ULL << component_id)) != 0;
}

static bool BECS_SignatureMatches(BECS_Signature entitySig, BECS_Signature querySig)
{
    return (entitySig & querySig) == querySig;
}

/* Memory Arena */

BECS_Arena BECS_CreateArena(void *memory, size_t capacity)
{
    assert(memory != NULL);
    assert(capacity > 0);

    /* Align base offset to 64 byte boundary */
    uintptr_t baseOffset = 64 - ((uintptr_t)memory % 64);
    BECS_Arena arena = {};
    arena.nextAllocation = baseOffset;
    arena.capacity = capacity;
    arena.memory = memory;
    return arena;
}

static void *BECS_AllocateMemory(BECS_Arena *arena, size_t size)
{
    assert(arena != NULL);
    assert(size > 0);

    size_t totalSizeBytes = size;
    uintptr_t nextAllocOffset = arena->nextAllocation + ((64 - (arena->nextAllocation % 64)) & 63);
    if (nextAllocOffset + totalSizeBytes <= arena->capacity)
    {
        arena->nextAllocation = nextAllocOffset + totalSizeBytes;
        return (void *)((uintptr_t)arena->memory + (uintptr_t)nextAllocOffset);
    }
    else
    {
        return NULL;
    }
}

/* Component Pool Backend */

static void BECS_ComponentPoolInit
(BECS_Arena *arena, BECS_ComponentPool *pool ,size_t componentSize, size_t capacity)
{
    assert(arena != NULL);
    assert(capacity <= BECS_MAX_ENTITIES);

    pool->unitSize = componentSize;
    pool->capacity = capacity;
    pool->count = 0;

    pool->denseArray = BECS_AllocateMemory(arena, componentSize * capacity);
    pool->denseEntityMapping = (BECS_Entity *)BECS_AllocateMemory(arena, sizeof(BECS_Entity) * capacity);

    uint32_t i;
    for (i = 0; i < BECS_MAX_ENTITIES; i++)
    {
        pool->sparseArray[i] = UINT32_MAX;
    }
}

static void *BECS_ComponentPoolAdd(BECS_ComponentPool *pool, BECS_Entity entity)
{
    assert(pool->count < pool->capacity);
    assert(entity < BECS_MAX_ENTITIES);

    uint32_t index = pool->count;
    pool->count++;
    pool->denseEntityMapping[index] = entity;
    pool->sparseArray[entity] = index;

    return (uint8_t *)pool->denseArray + (index * pool->unitSize);
}

static void BECS_ComponentPoolRemove(BECS_ComponentPool *pool, BECS_Entity entity)
{
    assert(entity < BECS_MAX_ENTITIES);

    uint32_t index = pool->sparseArray[entity];
    if (index == UINT32_MAX)
    {
        return;
    }

    uint32_t last_index = pool->count - 1;
    BECS_Entity last_entity = pool->denseEntityMapping[last_index];

    if (index != last_index)
    {
        void *dst = (uint8_t *)pool->denseArray + (index * pool->unitSize);
        void *src = (uint8_t *)pool->denseArray + (last_index * pool->unitSize);
        memcpy(dst, src, pool->unitSize);

        pool->denseEntityMapping[index] = last_entity;
        pool->sparseArray[last_entity] = index;
        pool->count--;
    }
    else
    {
        pool->count--;
    }

    pool->sparseArray[entity] = UINT32_MAX;
}

static void *BECS_ComponentPoolGet(BECS_ComponentPool *pool, BECS_Entity entity)
{
    assert(entity < BECS_MAX_ENTITIES);

    uint32_t index = pool->sparseArray[entity];
    if (index == UINT32_MAX)
    {
        return NULL;
    }

    return (uint8_t *)pool->denseArray + (index * pool->unitSize);
}

static bool BECS_ComponentPoolHas(BECS_ComponentPool *pool, BECS_Entity entity)
{
    assert(entity < BECS_MAX_ENTITIES);
    return pool->sparseArray[entity] != UINT32_MAX;
}

/* Helper API Functions */

size_t BECS_GetMinMemoryArenaSize(size_t *componentSizes, size_t arrayLength)
{
    assert(componentSizes != NULL);
    assert(arrayLength <= BECS_MAX_COMPONENTS);

    size_t totalSize = 0;
    totalSize += sizeof(BECS_ECS);
    totalSize += sizeof(BECS_ComponentPool) * BECS_MAX_COMPONENTS;
    totalSize += sizeof(BECS_Entity) * BECS_MAX_ENTITIES * BECS_MAX_COMPONENTS;

    size_t sumCompSizes = 0;
    uint32_t i;
    for (i = 0; i < arrayLength; i++)
    {
        sumCompSizes += componentSizes[i];
    }

    sumCompSizes *= BECS_MAX_ENTITIES;
    totalSize += sumCompSizes;

    return totalSize;
}

/* World API Functions */

BECS_ECS BECS_InitECS(void *memory, size_t size)
{
    assert(memory != NULL);
    assert(size > 0);

    BECS_ECS ecs;
    ecs.arena = BECS_CreateArena(memory, size);
    ecs.entityCount = 0;
    memset(ecs.entitySignatures, 0, sizeof(BECS_Signature) * BECS_MAX_ENTITIES);
    memset(ecs.aliveEntities, false, sizeof(bool) * BECS_MAX_ENTITIES);
    memset(ecs.freeEntities, BECS_NULL_ENTITY, sizeof(BECS_Entity) * BECS_MAX_ENTITIES);
    ecs.freeEntityCount = 0;
    ecs.nextEntity = 0;

    ecs.componentTypeCount = 0;
    uint32_t i;
    for (i = 0; i < BECS_MAX_COMPONENTS; i++)
    {
        ecs.componentPools[i].denseArray = NULL;
        ecs.componentPools[i].denseEntityMapping = NULL;
    }

    return ecs;
}

void BECS_ResetECS(BECS_ECS *ecs)
{
    if (!ecs)
    {
        return;
    }

    ecs->arena = BECS_CreateArena(ecs->arena.memory, ecs->arena.capacity);
    ecs->entityCount = 0;
    memset(ecs->entitySignatures, 0, sizeof(BECS_Signature) * BECS_MAX_ENTITIES);
    memset(ecs->aliveEntities, false, sizeof(bool) * BECS_MAX_ENTITIES);
    memset(ecs->freeEntities, BECS_NULL_ENTITY, sizeof(BECS_Entity) * BECS_MAX_ENTITIES);
    ecs->freeEntityCount = 0;
    ecs->nextEntity = 0;

    ecs->componentTypeCount = 0;
    uint32_t i;
    for (i = 0; i < BECS_MAX_COMPONENTS; i++)
    {
        ecs->componentPools[i].denseArray = NULL;
        ecs->componentPools[i].denseEntityMapping = NULL;
    }
}

/* Entity API Functions */

BECS_Entity BECS_EntityCreate(BECS_ECS *ecs)
{
    assert(ecs != NULL);

    BECS_Entity entity = BECS_NULL_ENTITY;

    if (ecs->freeEntityCount)
    {
        /* Pop entity off the back of the free list */
        entity = ecs->freeEntities[--ecs->freeEntityCount];
        ecs->entitySignatures[entity] = 0;
        ecs->aliveEntities[entity] = true;
        ecs->entityCount++;
    }
    else if (ecs->nextEntity < BECS_MAX_ENTITIES)
    {
        entity = ecs->nextEntity;
        ecs->aliveEntities[entity] = true;
        ecs->entitySignatures[entity] = 0;
        ecs->nextEntity++;
        ecs->entityCount++;
    }
    return entity;
}

void BECS_EntityDestroy(BECS_ECS *ecs, BECS_Entity entity)
{
    assert(ecs != NULL);
    assert(entity >= 0 && entity < BECS_MAX_ENTITIES);

    if (!ecs->aliveEntities[entity])
    {
        return;
    }

    uint32_t i;
    for (i = 0; i < BECS_MAX_COMPONENTS; i++)
    {
        BECS_Signature sig = ecs->entitySignatures[entity];
        if (BECS_SignatureHas(sig, i))
        {
            BECS_ComponentRemove(ecs, entity, i);
        }
    }

    ecs->aliveEntities[entity] = false;
    ecs->entitySignatures[entity] = 0;
    ecs->freeEntities[ecs->freeEntityCount] = entity;
    ecs->entityCount--;
    ecs->freeEntityCount++;
}

bool BECS_IsEntityAlive(BECS_ECS *ecs, BECS_Entity entity)
{
    assert(ecs != NULL);

    if (entity >= BECS_MAX_ENTITIES || entity == BECS_NULL_ENTITY)
    {
        return false;
    }
    return ecs->aliveEntities[entity];
}

/* Component API Functions */

BECS_ComponentID BECS_ComponentRegister(BECS_ECS *ecs, size_t componentSize, size_t maximumComponentCount)
{
    assert(ecs != NULL);
    assert(ecs->componentTypeCount < BECS_MAX_COMPONENTS);
    assert(componentSize > 0);

    if (maximumComponentCount == 0)
    {
        maximumComponentCount = BECS_MAX_ENTITIES;
    }

    BECS_ComponentID id = ecs->componentTypeCount;
    BECS_ComponentPoolInit(&ecs->arena, &ecs->componentPools[id],
            componentSize, maximumComponentCount);
    ecs->componentTypeCount++;

    return id;
}

void *BECS_ComponentAdd(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId)
{
    assert(ecs != NULL);
    assert(entity >= 0 && entity < BECS_MAX_ENTITIES);
    assert(componentId < ecs->componentTypeCount);
    assert(ecs->aliveEntities[entity]);

    void *component = BECS_ComponentPoolAdd(&ecs->componentPools[componentId], entity);
    ecs->entitySignatures[entity] = BECS_SignatureSet(ecs->entitySignatures[entity], componentId);
    memset(component, 0, ecs->componentPools[componentId].unitSize);
    return component;
}

void BECS_ComponentRemove(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId)
{
    assert(ecs != NULL);
    assert(entity >= 0 && entity < BECS_MAX_ENTITIES);
    assert(componentId < ecs->componentTypeCount);

    if (!ecs->aliveEntities[entity])
    {
        return;
    }
    BECS_ComponentPoolRemove(&ecs->componentPools[componentId], entity);
    ecs->entitySignatures[entity] &= ~(1ULL << componentId);
}

void *BECS_ComponentGet(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId)
{
    assert(ecs != NULL);
    assert(entity >= 0 && entity < BECS_MAX_ENTITIES);
    assert(componentId < ecs->componentTypeCount);

    if (!ecs->aliveEntities[entity])
    {
        return NULL;
    }
    return BECS_ComponentPoolGet(&ecs->componentPools[componentId], entity);
}

bool BECS_ComponentHas(BECS_ECS *ecs, BECS_Entity entity, BECS_ComponentID componentId)
{
    assert(ecs != NULL);
    assert(entity >= 0 && entity < BECS_MAX_ENTITIES);
    assert(componentId < ecs->componentTypeCount);

    if (!ecs->aliveEntities[entity])
    {
        return false;
    }
    return BECS_SignatureHas(ecs->entitySignatures[entity], componentId);
}

#endif /* BECS_IMPLEMENTATION */
#endif /* BECS_H */