#include "entity_map.h"

#include <math.h>

EntityMap entityMapCreate(const EntityMapDesc *desc) {
    i32 width = (desc->maxWidth + (desc->cellResolution - 1)) / desc->cellResolution;
    i32 height = (desc->maxHeight + (desc->cellResolution - 1)) / desc->cellResolution;
    EntityMap map = {
            .width=width,
            .height=height,
            .cellResolution=desc->cellResolution,
            .cellDepth=desc->cellDepth,
    };

    size_t numBytes = sizeof(EntityMapCell) * map.width * map.height * map.cellDepth;
    map.cells = bzAlloc(numBytes);
    bzMemSet(map.cells, 0, numBytes);
    map.entities = bzArrayNew(EntityMapEntry, 512);

    return map;
}
void entityMapDestroy(EntityMap *entityMap) {
    bzFree(entityMap->cells);
    entityMap->cells = NULL;
    bzArrayFree(entityMap->entities);
    entityMap->entities = NULL;
}

static EntityMapIndex calculateMapIndex(EntityMap *entityMap, Position pos, Size size) {
    i32 minX = (i32) floorf(pos.x);
    i32 minY = (i32) floorf(pos.y);
    i32 maxX = (i32) floorf(pos.x + size.x);
    i32 maxY = (i32) floorf(pos.y + size.y);
    return (EntityMapIndex) {
            (i16) (minX / entityMap->width),
            (i16) (minY / entityMap->height),
            (i16) (maxX / entityMap->width),
            (i16) (maxY / entityMap->height),
    };
}

static bool insertMapEntry(EntityMap *entityMap, EntityMapCell idx, i32 x, i32 y) {
    EntityMapCell *cell = entityMap->cells + (y * entityMap->width + x) * entityMap->cellDepth;
    for (i32 i = 0; i < entityMap->cellResolution; i++) {
        if (cell[i] == 0) {
            cell[i] = idx;
            return true;
        }
    }
    return false;
}
static void removeMapEntry(EntityMap *entityMap, EntityMapCell idx, i32 x, i32 y) {
    EntityMapCell *cell = entityMap->cells + (y * entityMap->width + x) * entityMap->cellDepth;
    i32 i;
    for (i = 0; i < entityMap->cellResolution; i++) {
        if (cell[i] == idx) {
            break;
        }
    }
    i32 j;
    for (j = i; j < entityMap->cellResolution; j++) {
        if (cell[j] == 0) {
            j--;
            break;
        }

    }
    if (i != j) {
        // swap with last
        cell[i] = cell[j];
        cell[j] = 0;
    } else {
        // 'i' is last
        cell[i] = 0;
    }
}

EntityMapEntry entityMapInsert(EntityMap *entityMap, ecs_entity_t entity, Position pos, Size size) {
    EntityMapEntry entry = {entity};
    EntityMapIndex mapIndex = calculateMapIndex(entityMap, pos, size);
    entry.index = mapIndex;

    i32 entryIdx = bzArraySize(entityMap->entities);
    entry.entryIdx = entryIdx;
    bzArrayPush(entityMap->entities, entry);

    bool wasInserted = true;

    for (i32 y = mapIndex.minY; y <= mapIndex.maxY; y++) {
        for (i32 x = mapIndex.minX; x <= mapIndex.maxX; x++) {
            wasInserted &= insertMapEntry(entityMap, entryIdx, x, y);
            if (!wasInserted) {
                entityMapRemoveIdx(entityMap, entryIdx);
                return (EntityMapEntry) {
                    .entity = 0,
                };
            }
        }
    }

    return entry;
}
EntityMapEntry entityMapUpdate(EntityMap *entityMap, EntityMapEntry entry, Position pos, Size size) {
    EntityMapEntry check = bzArrayGet(entityMap->entities, entry.entity);
    BZ_ASSERT(check.entity == entry.entity);

    EntityMapIndex mapIndex = calculateMapIndex(entityMap, pos, size);
    if (entry.index.minX == mapIndex.minX &&
        entry.index.minY == mapIndex.minY &&
        entry.index.maxX == mapIndex.maxX &&
        entry.index.maxY == mapIndex.maxY) {
        return entry;
    }

    // TODO: Has room for optimization
    entityMapRemoveIdx(entityMap, entry.entryIdx);
    entityMapInsert(entityMap, entry.entity, pos, size);

}
void entityMapRemove(EntityMap *entityMap, EntityMapEntry entry) {
    entityMapRemoveIdx(entityMap, entry.entryIdx);
}
void entityMapRemoveIdx(EntityMap *entityMap, i32 entryID) {
    EntityMapEntry entry = bzArrayGet(entityMap->entities, entryID);
    EntityMapIndex index = entry.index;

    for (i32 y = index.minY; y <= index.maxY; y++) {
       for (i32 x = index.minX; x <= index.maxX; x++) {
           removeMapEntry(entityMap, entryID, x, y);
       }
    }

    EntityMapEntry last = bzArrayGet(entityMap->entities, bzArraySize(entityMap->entities));
    bzArraySet(entityMap->entities, entryID, last);
    bzArrayPop(entityMap->entities);
}

EntityMapIter entityMapQueryIter(EntityMap *entityMap, Position pos, Size size) {
    EntityMapIndex mapIndex = calculateMapIndex(entityMap, pos, size);
    EntityMapIter it = {
            .entityMap=entityMap,
            .queryIdx=entityMap->queryCount++,
            .index=mapIndex,
            .x=mapIndex.minX,
            .y=mapIndex.minY,
            .cellIdx=0,
    };
    return it;
}
bool entityMapQueryNext(EntityMapIter *it) {
    i32 cellOffset = (it->y * it->entityMap->width + it->x) *
                     it->entityMap->cellDepth + it->cellIdx;
    if (it->cellIdx >= it->entityMap->cellDepth ||
        it->entityMap->cells[cellOffset] == 0) {
        it->cellIdx = 0;
        it->x++;
        if (it->x > it->index.maxX) {
            it->y++;
            if (it->y > it->index.maxY)
                return false;
        }
    }
    it->cellIdx++;

    EntityMapCell cell = it->entityMap->cells[cellOffset];
    EntityMapEntry *entry = &bzArrayGet(it->entityMap->entities, cell);
    if (entry->queryIdx == it->queryIdx) {
        return entityMapQueryNext(it);
    }
    entry->queryIdx = it->queryIdx;
    it->entry = *entry;

    return true;
}