#include "spatial_grid.h"

#include "../memory/memory.h"
#include "array.h"
#include "object_pool.h"

#include <math.h>

typedef struct BzSpatialGrid {
    BzSpatialGridID **cells;
    BzObjectPool *entriesPool;
    i32 width;
    i32 height;
    i32 cellWidth;
    i32 cellHeight;
    i32 userDataSize;
    i32 queryCount;
} BzSpatialGrid;

static BzSpatialGridID **getCell(const BzSpatialGrid *grid, i32 x, i32 y);

BzSpatialGrid *bzSpatialGridCreate(const BzSpatialGridDesc *desc) {
    i32 width = (desc->maxWidth + desc->cellWidth - 1) / desc->cellWidth;
    i32 height = (desc->maxHeight + desc->cellHeight - 1) / desc->cellHeight;

    size_t cellsNumBytes = width * height * sizeof(BzSpatialGridID *);
    size_t numBytes = sizeof(BzSpatialGrid) + cellsNumBytes;
    BzSpatialGrid *grid = bzAlloc(numBytes);

    grid->width = width;
    grid->height = height;
    grid->cellWidth = desc->cellWidth;
    grid->cellHeight = desc->cellHeight;
    grid->userDataSize = desc->userDataSize;
    grid->queryCount = 0;
    grid->cells = (BzSpatialGridID **) (grid + 1);
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            BzSpatialGridID **cell = getCell(grid, x, y);
            *cell = bzArrayCreate(BzSpatialGridID, 4);
        }
    }
    grid->entriesPool = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectSize=sizeof(BzSpatialGridEntry) + desc->userDataSize,
        .objectsPerPage=1024,
    });
    return grid;
}

void bzSpatialGridDestroy(BzSpatialGrid *grid) {
    bzObjectPoolDestroy(grid->entriesPool);
    grid->entriesPool = NULL;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            BzSpatialGridID **cell = getCell(grid, x, y);
            bzArrayDestroy(*cell);
            *cell = NULL;
        }
    }

    grid->cells = NULL;
    bzFree(grid);
}

void *bzSpatialGridGetData(const BzSpatialGrid *grid, BzSpatialGridID id) {
    return (u8 *)bzObjectPoolGetObject(grid->entriesPool, id) + sizeof(BzSpatialGridEntry);
}


static BzSpatialGridID **getCell(const BzSpatialGrid *grid, i32 x, i32 y) {
    return grid->cells + y * grid->width + x;
}

static void insertCell(BzSpatialGrid *grid, BzSpatialGridID id, i32 x, i32 y) {
    BzSpatialGridID **cell = getCell(grid, x, y);
    bzArrayPush(*cell, id);
}
static void removeCell(BzSpatialGrid *grid, BzSpatialGridID id, i32 x, i32 y) {
    BzSpatialGridID **cell = getCell(grid, x, y);
    i32 i;
    for (i = 0; i < bzArraySize(*cell); i++) {
        if (id == (*cell)[i]) {
            break;
        }
    }
    if (i < bzArraySize(*cell)) {
        BzSpatialGridID lastID = bzArrayPop(*cell);
        (*cell)[i] = lastID;
    }
}

static BzSpatialGridIndex calculateGridIndex(BzSpatialGrid *grid, f32 posX, f32 posY, f32 sizeX, f32 sizeY) {
    i32 minX = (i32) floorf(posX);
    i32 minY = (i32) floorf(posY);
    i32 maxX = (i32) ceilf(posX + sizeX);
    i32 maxY = (i32) ceilf(posY + sizeY);
    BzSpatialGridIndex index = {
            (i16) (minX / grid->cellWidth),
            (i16) (minY / grid->cellHeight),
            (i16) (maxX / grid->cellWidth),
            (i16) (maxY / grid->cellHeight),
    };
    if (index.minX < 0) index.minX = 0;
    if (index.minY < 0) index.minY = 0;
    if (index.maxX >= grid->width) index.maxX = grid->width - 1;
    if (index.maxY >= grid->height) index.maxY = grid->height - 1;

    if (index.minX > index.maxX) {
        i32 tmp = index.minX;
        index.minX = index.maxX;
        index.maxX = tmp;
    }
    if (index.minY > index.maxY) {
        i32 tmp = index.minY;
        index.minY = index.maxY;
        index.maxX = tmp;
    }
    return index;
}

BzSpatialGridID bzSpatialGridInsert(BzSpatialGrid *grid, void *data, f32 posX, f32 posY, f32 sizeX, f32 sizeY) {
    BzSpatialGridEntry *entry = bzObjectPool(grid->entriesPool);
    BZ_ASSERT(entry);

    BzSpatialGridIndex index = calculateGridIndex(grid, posX, posY, sizeX, sizeY);

    BzSpatialGridID id = bzObjectPoolGetIdx(grid->entriesPool, entry);

    for (i32 y = index.minY; y <= index.maxY; y++) {
        for (i32 x = index.minX; x <= index.maxX; x++) {
            insertCell(grid, id, x, y);
        }
    }

    entry->id = id;
    entry->queryIdx = 0;
    entry->index = index;
    if (grid->userDataSize) {
        bzMemCpy(entry + 1, data, grid->userDataSize);
    }

    return id;
}

static bool indexContains(BzSpatialGridIndex index, i32 x, i32 y) {
    return x >= index.minX && x <= index.maxX && y >= index.minY && y <= index.maxY;
}

void bzSpatialGridUpdate(BzSpatialGrid *grid, BzSpatialGridID id, f32 posX, f32 posY, f32 sizeX, f32 sizeY) {
    BzSpatialGridEntry *entry = bzObjectPoolGetObject(grid->entriesPool, id);
    BZ_ASSERT(entry && entry->id == id);

    BzSpatialGridIndex oldIndex = entry->index;
    BzSpatialGridIndex newIndex = calculateGridIndex(grid, posX, posY, sizeX, sizeY);
    if (oldIndex.minX == newIndex.minX &&
        oldIndex.minY == newIndex.minY &&
        oldIndex.maxX == newIndex.maxX &&
        oldIndex.maxY == newIndex.maxY)  {
        return;
    }

    for (i32 y = oldIndex.minY; y <= oldIndex.maxY; y++) {
        for (i32 x = oldIndex.minX; x <= oldIndex.maxX; x++) {
            bool inNew = indexContains(newIndex, x, y);
            if (!inNew) removeCell(grid, id, x, y);
        }
    }

    for (i32 y = newIndex.minY; y <= newIndex.maxY; y++) {
        for (i32 x = newIndex.minX; x <= newIndex.maxX; x++) {
            bool inOld = indexContains(oldIndex, x, y);
            if (!inOld) insertCell(grid, id, x, y);

        }
    }

    entry->index = newIndex;
}
void bzSpatialGridRemove(BzSpatialGrid *grid, BzSpatialGridID id) {
    BzSpatialGridEntry *entry = bzObjectPoolGetObject(grid->entriesPool, id);
    BZ_ASSERT(entry && entry->id == id);

    BzSpatialGridIndex index = entry->index;
    for (i32 y = index.minY; y <= index.maxY; y++) {
        for (i32 x = index.minX; x <= index.maxX; x++) {
            removeCell(grid, id, x, y);
        }
    }
    bzObjectPoolRelease(grid->entriesPool, entry);
}

BzSpatialGridIter bzSpatialGridIter(BzSpatialGrid *grid, f32 posX, f32 posY, f32 sizeX, f32 sizeY) {
    BzSpatialGridIndex index = calculateGridIndex(grid, posX, posY, sizeX, sizeY);
    BzSpatialGridIter it = {
            .index=index,
            ._grid=grid,
            ._queryIdx=grid->queryCount++,
            ._x=index.minX,
            ._y=index.minY,
            ._cellIdx=0

    };
    return it;
}
bool bzSpatialGridQueryNext(BzSpatialGridIter *it) {
    BzSpatialGridID **cell = getCell(it->_grid, it->_x, it->_y);

    BzSpatialGridID id;
    BzSpatialGridEntry *entry;
    bool foundUnprocessed = false;
    while (!foundUnprocessed) {
        while (it->_cellIdx >= bzArraySize(*cell)) {
            it->_cellIdx = 0;
            it->_x++;
            if (it->_x > it->index.maxX) {
                it->_x = it->index.minX;
                it->_y++;
                if (it->_y > it->index.maxY) {
                    return false;
                }
            }
            cell = getCell(it->_grid, it->_x, it->_y);
        }
        id = bzArrayGet(*cell, it->_cellIdx);
        entry = bzObjectPoolGetObject(it->_grid->entriesPool, id);
        BZ_ASSERT(entry && id == entry->id);
        if (entry->queryIdx != it->_queryIdx) {
            entry->queryIdx = it->_queryIdx;
            foundUnprocessed = true;
        }
        it->_cellIdx++;
    }

    it->entry = *entry;
    it->data = (void *) (entry + 1);

    return true;
}

#include <raylib.h>
#include <stdio.h>

void bzSpatialGridDrawDebugGrid(const BzSpatialGrid *grid) {
    // Draw lines
    i32 posX = 0;
    i32 posY = 0;
    for (i32 x = 0; x <= grid->width; x++) {
        DrawLine(posX, 0, posX, grid->height * grid->cellHeight, BLACK);
        posX += grid->cellWidth;
    }
    for (i32 y = 0; y <= grid->height; y++) {
        DrawLine(0, posY, grid->width * grid->cellWidth, posY, BLACK);
        posY += grid->cellHeight;
    }

    posX = 0;
    posY = 0;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            BzSpatialGridID **cell = getCell(grid, x, y);
            i32 numEntries = 0;
            if (*cell)
                numEntries = bzArraySize(*cell);
            Color color = WHITE;
            i32 tint = numEntries;
            if (tint > 10) tint = 10;
            tint = 10 - tint;
            color.g = 25 * tint;
            color.b = 25 * tint;
            char buf[8];
            snprintf(buf, sizeof(buf), "%d", numEntries);
            DrawText(buf, posX + 1, posY + 1, 6, color);

            posX += grid->cellWidth;
        }
        posX = 0;
        posY += grid->cellHeight;
    }

}