#include "space.h"
#include "math.h"
#include "arena.h"
#include "collider.h"

/* 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(struct space) + (sizeof(struct space) % alignof(struct space_entry)))

/* Accessed via sim_ent_nil() */
READONLY struct space_entry _g_space_entry_nil = { .valid = false };
READONLY struct space_cell _g_space_cell_nil = { .valid = false };
READONLY struct space _g_space_nil = { .valid = false };

/* ========================== *
 * 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. */
struct space *space_alloc(f32 cell_size, u32 num_bins_sqrt)
{
    struct space *space;
    {
        struct arena arena = arena_alloc(GIGABYTE(64));
        space = arena_push(&arena, struct space);
        space->entry_arena = arena;
    }

    space->valid = true;
    space->entries = arena_dry_push(&space->entry_arena, struct space_entry);

    space->cell_arena = arena_alloc(GIGABYTE(64));
    space->cell_size = cell_size;
    space->num_bins = num_bins_sqrt * num_bins_sqrt;
    space->num_bins_sqrt = num_bins_sqrt;
    space->bins = arena_push_array(&space->cell_arena, struct space_cell_bin, space->num_bins);

    return space;
}

void space_release(struct space *space)
{
    arena_release(&space->cell_arena);
    arena_release(&space->entry_arena);
}

void space_reset(struct space *space)
{
    arena_pop_to(&space->entry_arena, (u64)space->entries - (u64)space->entry_arena.base);
    arena_reset(&space->cell_arena);
    space->bins = arena_push_array(&space->cell_arena, struct space_cell_bin, space->num_bins);
    space->num_entries_reserved = 0;
    space->first_free_cell = NULL;
    space->first_free_cell_node = NULL;
    space->first_free_entry = NULL;
}

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

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

INTERNAL struct v2i32 world_to_cell_coords(f32 cell_size, struct v2 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 V2I32((i32)x, (i32)y);
}

INTERNAL i32 cell_coords_to_bin_index(struct space *space, struct v2i32 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;
}

struct space_cell *space_get_cell(struct space *space, struct v2i32 cell_pos)
{
    i32 bin_index = cell_coords_to_bin_index(space, cell_pos);
    struct space_cell_bin *bin = &space->bins[bin_index];
    struct space_cell *res = space_cell_nil();
    for (struct space_cell *n = bin->first_cell; n; n = n->next_in_bin) {
        if (v2i32_eq(n->pos, cell_pos)) {
            res = n;
            break;
        }
    }
    return res;
}

INTERNAL void space_cell_node_alloc(struct v2i32 cell_pos, struct space_entry *entry)
{
    struct space *space = space_from_entry(entry);
    i32 bin_index = cell_coords_to_bin_index(space, cell_pos);
    struct space_cell_bin *bin = &space->bins[bin_index];

    /* Find existing cell */
    struct space_cell *cell = NULL;
    for (struct space_cell *n = bin->first_cell; n; n = n->next_in_bin) {
        if (v2i32_eq(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 = arena_push_no_zero(&space->cell_arena, struct space_cell);
        }
        MEMZERO_STRUCT(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 = true;
    }

    /* Allocate node */
    struct space_cell_node *node;
    {
        if (space->first_free_cell_node) {
            node = space->first_free_cell_node;
            space->first_free_cell_node = node->next_free;
        } else {
            node = arena_push_no_zero(&space->cell_arena, struct space_cell_node);
        }
        MEMZERO_STRUCT(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(struct space_cell_node *n)
{
    struct space_cell *cell = n->cell;
    struct space_entry *entry = n->entry;
    struct space *space = space_from_entry(entry);
    struct space_cell_bin *bin = cell->bin;

    /* Remove from entry list */
    {
        struct space_cell_node *prev = n->prev_in_entry;
        struct space_cell_node *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 */
    {
        struct space_cell_node *prev = n->prev_in_cell;
        struct space_cell_node *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 */
        struct space_cell *prev = cell->prev_in_bin;
        struct space_cell *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 = false;

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

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

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

struct space_entry *space_entry_from_handle(struct space *space, struct space_entry_handle handle)
{
    struct space_entry *entry = space_entry_nil();

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

    return entry;
}

struct space_entry *space_entry_alloc(struct space *space, struct sim_ent_id ent)
{
    struct space_entry *entry = NULL;
    struct space_entry_handle handle = ZI;
    if (space->first_free_entry) {
        entry = space->first_free_entry;
        space->first_free_entry = entry->next_free;
        handle = entry->handle;
    } else {
        entry = arena_push_no_zero(&space->entry_arena, struct space_entry);
        handle.idx = space->num_entries_reserved;
        handle.gen = 1;
        ++space->num_entries_reserved;
    }
    MEMZERO_STRUCT(entry);
    entry->valid = true;
    entry->handle = handle;
    entry->ent = ent;
    return entry;
}

void space_entry_release(struct space_entry *entry)
{
    /* Release nodes */
    struct space_cell_node *n = entry->first_node;
    while (n) {
        struct space_cell_node *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;
    }

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

void space_entry_update_aabb(struct space_entry *entry, struct aabb new_aabb)
{
    struct space *space = space_from_entry(entry);
    f32 cell_size = space->cell_size;

    struct v2i32 old_cell_p0 = V2I32(0, 0);
    struct v2i32 old_cell_p1 = V2I32(0, 0);
    if (entry->first_node) {
        struct 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);
    }

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

    /* Release outdated nodes */
    struct space_cell_node *n = entry->first_node;
    while (n) {
        struct space_cell *cell = n->cell;
        struct v2i32 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 */
            struct space_cell_node *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(V2I32(x, y), entry);
            }
        }
    }

    entry->aabb = new_aabb;
}

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

struct space_iter space_iter_begin_aabb(struct space *space, struct aabb aabb)
{
    struct space_iter 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 */
        struct v2i32 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;
}

struct space_entry *space_iter_next(struct space_iter *iter)
{
    struct space *space = iter->space;
    struct aabb iter_aabb = iter->aabb;
    struct v2i32 cell_start = iter->cell_start;
    struct v2i32 cell_end = iter->cell_end;
    struct v2i32 cell_cur = iter->cell_cur;
    i32 span = cell_end.x - cell_start.x;

    struct space_cell_node *next_node = NULL;
    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 != NULL);
        next_node = iter->prev->next_in_cell;
    } else if (cell_cur.x > cell_end.x || cell_cur.y > cell_end.y) {
        /* Ended */
        return NULL;
    }

    while (true) {
        if (next_node) {
            struct space_entry *entry = next_node->entry;
            struct aabb entry_aabb = entry->aabb;
            if (collider_test_aabb(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);
                struct space_cell *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 : NULL;
}