/* FIXME: Default space entry & cell pointers to nil */

/* Offset in bytes from start of space struct to start of entry array (assume adjacently allocated) */
#define SPACE_ENTRIES_OFFSET (sizeof(Space) + (sizeof(Space) % alignof(SpaceEntry)))

/* Accessed via sim_ent_nil() */
Readonly SpaceEntry _g_space_entry_nil = { .valid = 0 };
Readonly SpaceCell _g_space_cell_nil = { .valid = 0 };
Readonly Space _g_space_nil = { .valid = 0 };

/* ========================== *
 * Space alloc
 * ========================== */

/* NOTE:
 * The number of bins determines how often tiles will collide in the spatial hash.
 * For example, at `num_bins_sqrt` = 256 (65536 bins), tiles <1, 1>, <1, 257>, and <257, 257> will collide. */
Space *space_alloc(f32 cell_size, u32 num_bins_sqrt)
{
    Space *space;
    {
        Arena *arena = AllocArena(Gibi(64));
        space = PushStruct(arena, Space);
        space->entry_arena = arena;
    }

    space->valid = 1;
    space->entries = PushDry(space->entry_arena, SpaceEntry);

    space->cell_arena = AllocArena(Gibi(64));
    space->cell_size = cell_size;
    space->num_bins = num_bins_sqrt * num_bins_sqrt;
    space->num_bins_sqrt = num_bins_sqrt;
    space->bins = PushStructs(space->cell_arena, SpaceCellBin, space->num_bins);

    return space;
}

void space_release(Space *space)
{
    ReleaseArena(space->cell_arena);
    ReleaseArena(space->entry_arena);
}

void space_reset(Space *space)
{
    PopTo(space->entry_arena, (u64)space->entries - (u64)ArenaBase(space->entry_arena));
    ResetArena(space->cell_arena);
    space->bins = PushStructs(space->cell_arena, SpaceCellBin, space->num_bins);
    space->num_entries_reserved = 0;
    space->first_free_cell = 0;
    space->first_free_cell_node = 0;
    space->first_free_entry = 0;
}

Space *space_from_entry(SpaceEntry *entry)
{
    if (entry->valid) {
        u64 first_entry_addr = (u64)(entry - entry->handle.idx);
        Space *space = (Space *)(first_entry_addr - SPACE_ENTRIES_OFFSET);
        Assert(space->entries == (SpaceEntry *)first_entry_addr);
        return space;
    } else {
        return space_nil();
    }
}

/* ========================== *
 * Cell
 * ========================== */

internal Vec2I32 world_to_cell_coords(f32 cell_size, Vec2 world_pos)
{
    f32 x = world_pos.x;
    f32 y = world_pos.y;
    x = (x + ((x >= 0) - (x < 0)) * cell_size) / cell_size;
    y = (y + ((y >= 0) - (y < 0)) * cell_size) / cell_size;
    return VEC2I32((i32)x, (i32)y);
}

internal i32 cell_coords_to_bin_index(Space *space, Vec2I32 cell_pos)
{
    i32 num_bins_sqrt = space->num_bins_sqrt;

    /* Cell pos of 0 is not valid and will be converted to -1 */
    Assert(cell_pos.x != 0 && cell_pos.y != 0);

    i32 index_x = cell_pos.x;
    i32 index_y = cell_pos.y;
    /* Offset cell index by -1 since cell pos of 0 is invalid */
    index_x -= (index_x >= 0);
    index_y -= (index_y >= 0);
    /* Un-mirror coords to prevent collisions between cells near the axes. (e.g. <3, 1> & <3, -1> should not collide) */
    index_x += (index_x < 0) * (num_bins_sqrt * ((index_x / -num_bins_sqrt) + 1));
    index_y += (index_y < 0) * (num_bins_sqrt * ((index_y / -num_bins_sqrt) + 1));

    i32 bin_index = (index_x % num_bins_sqrt) + (index_y % num_bins_sqrt) * num_bins_sqrt;
    Assert(bin_index >= 0 && bin_index < (i32)space->num_bins);

    return bin_index;
}

SpaceCell *space_get_cell(Space *space, Vec2I32 cell_pos)
{
    i32 bin_index = cell_coords_to_bin_index(space, cell_pos);
    SpaceCellBin *bin = &space->bins[bin_index];
    SpaceCell *result = space_cell_nil();
    for (SpaceCell *n = bin->first_cell; n; n = n->next_in_bin) {
        if (EqVec2I32(n->pos, cell_pos)) {
            result = n;
            break;
        }
    }
    return result;
}

internal void space_cell_node_alloc(Vec2I32 cell_pos, SpaceEntry *entry)
{
    Space *space = space_from_entry(entry);
    i32 bin_index = cell_coords_to_bin_index(space, cell_pos);
    SpaceCellBin *bin = &space->bins[bin_index];

    /* Find existing cell */
    SpaceCell *cell = 0;
    for (SpaceCell *n = bin->first_cell; n; n = n->next_in_bin) {
        if (EqVec2I32(n->pos, cell_pos)) {
            cell = n;
            break;
        }
    }

    /* Allocate new cell if necessary */
    if (!cell) {
        if (space->first_free_cell) {
            cell = space->first_free_cell;
            space->first_free_cell = cell->next_free;
        } else {
            cell = PushStructNoZero(space->cell_arena, SpaceCell);
        }
        ZeroStruct(cell);
        if (bin->last_cell) {
            bin->last_cell->next_in_bin = cell;
            cell->prev_in_bin = bin->last_cell;
        } else {
            bin->first_cell = cell;
        }
        bin->last_cell = cell;
        cell->pos = cell_pos;
        cell->bin = bin;
        cell->valid = 1;
    }

    /* Allocate node */
    SpaceCellNode *node;
    {
        if (space->first_free_cell_node) {
            node = space->first_free_cell_node;
            space->first_free_cell_node = node->next_free;
        } else {
            node = PushStructNoZero(space->cell_arena, SpaceCellNode);
        }
        ZeroStruct(node);
    }

    /* Insert into cell list */
    node->cell = cell;
    if (cell->last_node) {
        cell->last_node->next_in_cell = node;
        node->prev_in_cell = cell->last_node;
    } else {
        cell->first_node = node;
    }
    cell->last_node = node;

    /* Insert into entry list */
    node->entry = entry;
    if (entry->last_node) {
        entry->last_node->next_in_entry = node;
        node->prev_in_entry = entry->last_node;
    } else {
        entry->first_node = node;
    }
    entry->last_node = node;
}

internal void space_cell_node_release(SpaceCellNode *n)
{
    SpaceCell *cell = n->cell;
    SpaceEntry *entry = n->entry;
    Space *space = space_from_entry(entry);
    SpaceCellBin *bin = cell->bin;

    /* Remove from entry list */
    {
        SpaceCellNode *prev = n->prev_in_entry;
        SpaceCellNode *next = n->next_in_entry;
        if (prev) {
            prev->next_in_entry = next;
        } else {
            entry->first_node = next;
        }
        if (next) {
            next->prev_in_entry = prev;
        } else {
            entry->last_node = prev;
        }
    }

    /* Remove from cell list */
    {
        SpaceCellNode *prev = n->prev_in_cell;
        SpaceCellNode *next = n->next_in_cell;
        if (prev) {
            prev->next_in_cell = next;
        } else {
            cell->first_node = next;
        }
        if (next) {
            next->prev_in_cell = prev;
        } else {
            cell->last_node = prev;
        }
    }

    /* If cell is now empty, release it */
    if (!cell->first_node && !cell->last_node) {
        /* Remove from bin */
        SpaceCell *prev = cell->prev_in_bin;
        SpaceCell *next = cell->next_in_bin;
        if (prev) {
            prev->next_in_bin = next;
        } else {
            bin->first_cell = next;
        }
        if (next) {
            next->prev_in_bin = prev;
        } else {
            bin->last_cell = prev;
        }
        cell->valid = 0;

        /* Insert cell into free list */
        cell->next_free = space->first_free_cell;
        space->first_free_cell = cell;
    }

    /* Insert node into free list */
    n->next_free = space->first_free_cell_node;
    space->first_free_cell_node = n;
}

/* ========================== *
 * Entry
 * ========================== */

SpaceEntry *space_entry_from_handle(Space *space, SpaceEntryHandle handle)
{
    SpaceEntry *entry = space_entry_nil();

    if (handle.gen > 0 && handle.idx < space->num_entries_reserved) {
        SpaceEntry *tmp = &space->entries[handle.idx];
        if (tmp->handle.gen == handle.gen) {
            entry = tmp;
        }
    }

    return entry;
}

SpaceEntry *space_entry_alloc(Space *space, EntityId ent)
{
    SpaceEntry *entry = 0;
    SpaceEntryHandle handle = ZI;
    if (space->first_free_entry) {
        entry = space->first_free_entry;
        space->first_free_entry = entry->next_free;
        handle = entry->handle;
    } else {
        entry = PushStructNoZero(space->entry_arena, SpaceEntry);
        handle.idx = space->num_entries_reserved;
        handle.gen = 1;
        ++space->num_entries_reserved;
    }
    ZeroStruct(entry);
    entry->valid = 1;
    entry->handle = handle;
    entry->ent = ent;
    return entry;
}

void space_entry_release(SpaceEntry *entry)
{
    /* Release nodes */
    SpaceCellNode *n = entry->first_node;
    while (n) {
        SpaceCellNode *next = n->next_in_entry;
        /* TODO: More efficient batch release that doesn't care about maintaining entry list */
        space_cell_node_release(n);
        n = next;
    }

    Space *space = space_from_entry(entry);
    entry->next_free = space->first_free_entry;
    entry->valid = 0;
    ++entry->handle.gen;
    space->first_free_entry = entry;
}

void space_entry_update_aabb(SpaceEntry *entry, Aabb new_aabb)
{
    Space *space = space_from_entry(entry);
    f32 cell_size = space->cell_size;

    Vec2I32 old_cell_p0 = VEC2I32(0, 0);
    Vec2I32 old_cell_p1 = VEC2I32(0, 0);
    if (entry->first_node) {
        Aabb old_aabb = entry->aabb;
        old_cell_p0 = world_to_cell_coords(cell_size, old_aabb.p0);
        old_cell_p1 = world_to_cell_coords(cell_size, old_aabb.p1);
    }

    Vec2I32 new_cell_p0 = world_to_cell_coords(cell_size, new_aabb.p0);
    Vec2I32 new_cell_p1 = world_to_cell_coords(cell_size, new_aabb.p1);

    /* Release outdated nodes */
    SpaceCellNode *n = entry->first_node;
    while (n) {
        SpaceCell *cell = n->cell;
        Vec2I32 cell_pos = cell->pos;
        if (cell_pos.x < new_cell_p0.x || cell_pos.x > new_cell_p1.x || cell_pos.y < new_cell_p0.y || cell_pos.y > new_cell_p1.y) {
            /* Cell is outside of new AABB */
            SpaceCellNode *next = n->next_in_entry;
            space_cell_node_release(n);
            n = next;
        } else {
            n = n->next_in_entry;
        }
    }

    /* Insert new nodes */
    for (i32 y = new_cell_p0.y; y <= new_cell_p1.y; ++y) {
        for (i32 x = new_cell_p0.x; x <= new_cell_p1.x; ++x) {
            if (x != 0 && y != 0 && (x < old_cell_p0.x || x > old_cell_p1.x || y < old_cell_p0.y || y > old_cell_p1.y)) {
                /* Cell is outside of old AABB */
                space_cell_node_alloc(VEC2I32(x, y), entry);
            }
        }
    }

    entry->aabb = new_aabb;
}

/* ========================== *
 * Iter
 * ========================== */

SpaceIter space_iter_begin_aabb(Space *space, Aabb aabb)
{
    SpaceIter iter = ZI;
    f32 cell_size = space->cell_size;

    iter.space = space;
    iter.cell_start = world_to_cell_coords(cell_size, aabb.p0);
    iter.cell_end = world_to_cell_coords(cell_size, aabb.p1);
    if (iter.cell_start.x > iter.cell_end.x || iter.cell_start.y > iter.cell_end.y) {
        /* Swap cell_start & cell_end */
        Vec2I32 tmp = iter.cell_start;
        iter.cell_start = iter.cell_end;
        iter.cell_end = tmp;
    }

    iter.aabb = aabb;
    iter.cell_cur = iter.cell_start;
    iter.cell_cur.x -= 1;
    iter.cell_cur.y -= 1;

    return iter;
}

SpaceEntry *space_iter_next(SpaceIter *iter)
{
    Space *space = iter->space;
    Aabb iter_aabb = iter->aabb;
    Vec2I32 cell_start = iter->cell_start;
    Vec2I32 cell_end = iter->cell_end;
    Vec2I32 cell_cur = iter->cell_cur;
    i32 span = cell_end.x - cell_start.x;

    SpaceCellNode *next_node = 0;
    if (cell_cur.x >= cell_start.x && cell_cur.x <= cell_end.x && cell_cur.y >= cell_start.y && cell_cur.y <= cell_end.y) {
        /* Started */
        Assert(iter->prev != 0);
        next_node = iter->prev->next_in_cell;
    } else if (cell_cur.x > cell_end.x || cell_cur.y > cell_end.y) {
        /* Ended */
        return 0;
    }

    for (;;) {
        if (next_node) {
            SpaceEntry *entry = next_node->entry;
            Aabb entry_aabb = entry->aabb;
            if (CLD_TestAabb(entry_aabb, iter_aabb)) {
                break;
            } else {
                next_node = next_node->next_in_cell;
            }
        } else {
            /* Reached end of cell, find next cell */
            b32 nextx = (cell_cur.x + 1) <= cell_end.x;
            b32 nexty = (nextx == 0) && ((cell_cur.y + 1) <= cell_end.y);
            if (nextx || nexty) {
                cell_cur.x += nextx - (span * nexty);
                cell_cur.y += nexty;
                cell_cur.x += (cell_cur.x == 0);
                cell_cur.y += (cell_cur.y == 0);
                SpaceCell *cell = space_get_cell(space, cell_cur);
                next_node = cell->first_node;
            } else {
                /* Reached end of iter */
                cell_cur.x += 1;
                cell_cur.y += 1;
                break;
            }
        }
    }
    iter->prev = next_node;
    iter->cell_cur = cell_cur;

    return next_node ? next_node->entry : 0;
}