#include "entity.h"
#include "math.h"

/* Offset in bytes from start of store struct to start of entities array (assume adjacently allocated) */
#define STORE_ENTITIES_OFFSET (sizeof(struct entity_store) + (sizeof(struct entity_store) % alignof(struct entity)))

/* Accessed via entity_store_nil() */
READONLY struct entity_store _g_entity_store_nil = { .valid = false };

/* Accessed via entity_nil() */
/* TODO: Allocate nil entity in nil store */
READONLY struct entity _g_entity_nil = {
    .valid = false,
    .local_xform = XFORM_IDENT_NOCAST,
    .cached_global_xform = XFORM_IDENT_NOCAST,
    .cached_global_xform_dirty = false,
    .friction = 0.5f,
    .mass_unscaled = 1,
    .inertia_unscaled = 1,
    .sprite_local_xform = XFORM_IDENT_NOCAST,
    .sprite_tint = COLOR_WHITE
};

/* ========================== *
 * Store allocation
 * ========================== */

INTERNAL struct entity *entity_alloc_internal(struct entity_store *store);

INTERNAL void store_make_root(struct entity_store *store)
{
    struct entity *root = entity_alloc_internal(store);
    root->is_root = true;
    root->local_xform = XFORM_IDENT;
    root->cached_global_xform = XFORM_IDENT;
    root->cached_global_xform_dirty = false;
    store->root = root->handle;
}

struct entity_store *entity_store_alloc(void)
{
    struct arena arena = arena_alloc(GIGABYTE(64));
    struct entity_store *store = arena_push_zero(&arena, struct entity_store);
    store->valid = true;
    store->arena = arena;
    store->entities = arena_dry_push(&arena, struct entity);
    ASSERT((u64)store->entities - (u64)store == STORE_ENTITIES_OFFSET);  /* Offset must be correct */
    store_make_root(store);
    return store;
}

void entity_store_release(struct entity_store *store)
{
    arena_release(&store->arena);
}

void entity_store_copy_replace(struct entity_store *dest, struct entity_store *src)
{
    struct arena dest_arena = dest->arena;
    struct entity *dest_entities = dest->entities;
    MEMCPY_STRUCT(dest, src);
    arena_copy_replace(&dest_arena, &src->arena);
    dest->arena = dest_arena;
    dest->entities = dest_entities;
}

/* ========================== *
 * Entity allocation
 * ========================== */

INTERNAL struct entity *entity_alloc_internal(struct entity_store *store)
{
    struct entity *entity = NULL;
    struct entity_handle handle = ZI;
    if (store->first_free.gen) {
        /* Reuse from free list */;
        entity = entity_from_handle(store, store->first_free);
        handle = entity->handle;
        ++handle.gen;
        store->first_free = entity->next_free;
    } else {
        /* Make new */
        entity = arena_push(&store->arena, struct entity);
        handle = (struct entity_handle) { .gen = 1, .idx = store->reserved++ };
    }
    *entity = _g_entity_nil;
    entity->valid = true;
    entity->handle = handle;
    entity->cached_global_xform_dirty = true;
    ++store->allocated;
    return entity;
}

struct entity *entity_alloc(struct entity *parent)
{
    ASSERT(parent->valid);
    struct entity_store *store = entity_store_from_entity(parent);
    struct entity *e = entity_alloc_internal(store);
    entity_link_parent(e, parent);
    return e;
}

INTERNAL void entity_release_internal(struct entity_store *store, struct entity *ent)
{
    /* Release children */
    struct entity_handle first_handle = ent->first;
    if (first_handle.gen) {
        for (struct entity *child = entity_from_handle(store, first_handle); child->valid; child = entity_from_handle(store, child->next)) {
            entity_release_internal(store, child);
        }
    }

    /* Release */
    ++ent->handle.gen;
    ent->valid = false;
    ent->next_free = store->first_free;
    store->first_free = ent->handle;
    --store->allocated;
}

void entity_release(struct entity_store *store, struct entity *ent)
{
    if (ent->parent.gen) {
        entity_unlink_from_parent(ent);
    }
    entity_release_internal(store, ent);
}

/* ========================== *
 * Query
 * ==========================  */

struct entity_store *entity_store_from_entity(struct entity *ent)
{
    if (ent->valid) {
        u64 first_entity_addr = (u64)(ent - ent->handle.idx);
        struct entity_store *store = (struct entity_store *)(first_entity_addr - STORE_ENTITIES_OFFSET);
        ASSERT(store->entities == (struct entity *)first_entity_addr);
        return store;
    } else {
        return entity_store_nil();
    }
}

/* Returns a valid entity or read-only nil entity. Always safe to read result, need to check `valid` to write. */
struct entity *entity_from_handle(struct entity_store *store, struct entity_handle handle)
{
    if (handle.gen != 0 && handle.idx < store->reserved) {
        struct entity *entity = &store->entities[handle.idx];
        if (entity->handle.gen == handle.gen) {
            return entity;
        }
    }
    return entity_nil();
}

struct entity *entity_find_first_match_one(struct entity_store *store, enum entity_prop prop)
{
    u64 count = store->reserved;
    struct entity *entities = store->entities;
    for (u64 entity_index = 0; entity_index < count; ++entity_index) {
        struct entity *ent = &entities[entity_index];
        if (ent->valid && entity_has_prop(ent, prop)) {
            return ent;
        }
    }
    return entity_nil();
}

struct entity *entity_find_first_match_all(struct entity_store *store, struct entity_prop_array props)
{
    u64 count = store->reserved;
    struct entity *entities = store->entities;
    for (u64 entity_index = 0; entity_index < count; ++entity_index) {
        struct entity *ent = &entities[entity_index];
        if (ent->valid) {
            b32 all = true;
            for (u64 i = 0; i < props.count; ++i) {
                if (!entity_has_prop(ent, props.props[i])) {
                    all = false;
                    break;
                }
            }
            if (all) {
                return ent;
            }
        }
    }
    return entity_nil();
}

/* ========================== *
 * Xform
 * ========================== */

INTERNAL void entity_mark_child_xforms_dirty(struct entity_store *store, struct entity *ent)
{
    for (struct entity *child = entity_from_handle(store, ent->first); child->valid; child = entity_from_handle(store, child->next)) {
        if (child->cached_global_xform_dirty) {
            break;
        } else {
            child->cached_global_xform_dirty = true;
            entity_mark_child_xforms_dirty(store, child);
        }
    }
}

INTERNAL struct xform entity_get_xform_w_store(struct entity_store *store, struct entity *ent)
{
    struct xform xf;
    if (ent->cached_global_xform_dirty) {
        if (ent->is_top) {
            xf = ent->local_xform;
        } else {
            struct entity *parent = entity_from_handle(store, ent->parent);
            xf = entity_get_xform_w_store(store, parent);
            xf = xform_mul(xf, ent->local_xform);
            ent->cached_global_xform = xf;
            ent->cached_global_xform_dirty = false;
        }
        ent->cached_global_xform = xf;
        ent->cached_global_xform_dirty = false;
    } else {
        xf = ent->cached_global_xform;
    }
    return xf;
}

struct xform entity_get_xform(struct entity *ent)
{
    struct xform xf;
    if (ent->cached_global_xform_dirty) {
        if (ent->is_top) {
            xf = ent->local_xform;
        } else {
            struct entity_store *store = entity_store_from_entity(ent);
            struct entity *parent = entity_from_handle(store, ent->parent);
            xf = entity_get_xform_w_store(store, parent);
            xf = xform_mul(xf, ent->local_xform);
            ent->cached_global_xform = xf;
            ent->cached_global_xform_dirty = false;
        }
        ent->cached_global_xform = xf;
        ent->cached_global_xform_dirty = false;
    } else {
        xf = ent->cached_global_xform;
    }
    return xf;
}

struct xform entity_get_local_xform(struct entity *ent)
{
    return ent->local_xform;
}

void entity_set_xform(struct entity *ent, struct xform xf)
{
    if (!xform_eq(xf, ent->cached_global_xform)) {
        struct entity_store *store = entity_store_from_entity(ent);
        /* Update local xform */
        if (ent->is_top) {
            ent->local_xform = xf;
        } else {
            struct entity *parent = entity_from_handle(store, ent->parent);
            struct xform parent_global = entity_get_xform_w_store(store, parent);
            ent->local_xform = xform_mul(xform_invert(parent_global), xf);
        }
        ent->cached_global_xform = xf;
        ent->cached_global_xform_dirty = false;
        entity_mark_child_xforms_dirty(store, ent);
    }
}

void entity_set_local_xform(struct entity *ent, struct xform xf)
{
    if (!xform_eq(xf, ent->local_xform)) {
        ent->local_xform = xf;
        ent->cached_global_xform_dirty = true;
        entity_mark_child_xforms_dirty(entity_store_from_entity(ent), ent);
    }
}

/* ========================== *
 * Movement
 * ========================== */

void entity_apply_linear_impulse(struct entity *ent, struct v2 impulse, struct v2 point)
{
    struct xform xf = entity_get_xform(ent);
    struct v2 center = xf.og;
    f32 scale = math_fabs(xform_get_determinant(xf));
    f32 inv_mass = 1.f / (ent->mass_unscaled * scale);
    f32 inv_inertia = 1.f / (ent->inertia_unscaled * scale);

    struct v2 vcp = v2_sub(point, center);
    entity_set_linear_velocity(ent, v2_add(ent->linear_velocity, v2_mul(impulse, inv_mass)));
    entity_set_angular_velocity(ent, v2_wedge(vcp, impulse) * inv_inertia);
}

void entity_apply_linear_impulse_to_center(struct entity *ent, struct v2 impulse)
{
    struct xform xf = entity_get_xform(ent);
    f32 scale = math_fabs(xform_get_determinant(xf));
    f32 inv_mass = 1.f / (ent->mass_unscaled * scale);

    entity_set_linear_velocity(ent, v2_add(ent->linear_velocity, v2_mul(impulse, inv_mass)));
}

void entity_apply_force_to_center(struct entity *ent, struct v2 force)
{
    ent->force = v2_add(ent->force, force);
}

void entity_apply_angular_impulse(struct entity *ent, f32 impulse)
{
    struct xform xf = entity_get_xform(ent);
    f32 scale = math_fabs(xform_get_determinant(xf));
    f32 inv_inertia = 1.f / (ent->inertia_unscaled * scale);
    entity_set_angular_velocity(ent, ent->angular_velocity + impulse * inv_inertia);
}

void entity_apply_torque(struct entity *ent, f32 torque)
{
    ent->torque += torque;
}

/* ========================== *
 * Tree
 * ========================== */

void entity_link_parent(struct entity *ent, struct entity *parent)
{
    struct entity_store *store = entity_store_from_entity(ent);

    if (ent->parent.gen) {
        /* Unlink from current parent */
        entity_unlink_from_parent(ent);
    }

    struct entity_handle handle = ent->handle;
    struct entity_handle parent_handle = parent->handle;

    ent->parent = parent_handle;

    struct entity_handle last_child_handle = parent->last;
    struct entity *last_child = entity_from_handle(store, last_child_handle);
    if (last_child->valid) {
        ent->prev = last_child_handle;
        last_child->next = handle;
    } else {
        parent->first = handle;
    }
    parent->last = handle;

    if (parent->is_root) {
        ent->is_top = true;
        ent->top = handle;
    } else {
        ent->top = parent->top;
    }
}

/* NOTE: Entity will be dangling after calling this, should re-link to root entity. */
void entity_unlink_from_parent(struct entity *ent)
{
    struct entity_store *store = entity_store_from_entity(ent);

    struct entity_handle parent_handle = ent->parent;
    struct entity *parent = entity_from_handle(store, parent_handle);
    struct entity *prev = entity_from_handle(store, ent->prev);
    struct entity *next = entity_from_handle(store, ent->next);

    /* Unlink from parent & siblings */
    if (prev->valid) {
        prev->next = next->handle;
    } else {
        parent->first = next->handle;
    }
    if (next->valid) {
        next->prev = prev->handle;
    } else {
        parent->last = prev->handle;
    }
    ent->prev = entity_handle_nil();
    ent->next = entity_handle_nil();
}

/* ========================== *
 * Entity lookup
 * ========================== */

struct entity_lookup entity_lookup_alloc(u64 num_buckets)
{
    ASSERT(num_buckets > 0);
    struct entity_lookup l = ZI;
    l.arena = arena_alloc(GIGABYTE(64));
    l.buckets = arena_push_array_zero(&l.arena, struct entity_lookup_bucket, num_buckets);
    l.num_buckets = num_buckets;
    return l;
}

void entity_lookup_release(struct entity_lookup *l)
{
    arena_release(&l->arena);
}

struct entity_lookup_entry *entity_lookup_get(struct entity_lookup *l, struct entity_lookup_key key)
{
    u64 index = key.hash % l->num_buckets;
    struct entity_lookup_bucket *bucket = &l->buckets[index];
    struct entity_lookup_entry *res = NULL;
    for (struct entity_lookup_entry *e = bucket->first; e; e = e->next) {
        if (e->key.hash == key.hash) {
            res = e;
            break;
        }
    }
    return res;
}

void entity_lookup_set(struct entity_lookup *l, struct entity_lookup_key key, struct entity_handle handle)
{
    u64 index = key.hash % l->num_buckets;
    struct entity_lookup_bucket *bucket = &l->buckets[index];

    struct entity_lookup_entry *prev = NULL;
    struct entity_lookup_entry **slot = &bucket->first;
    while (*slot) {
        if ((*slot)->key.hash == key.hash) {
            break;
        }
        prev = *slot;
        slot = &(*slot)->next;
    }

    struct entity_lookup_entry *entry = *slot;
    if (entry) {
        /* Set existing entry */
        entry->entity = handle;
    } else {
        /* Allocate entry */
        if (l->first_free_entry) {
            entry = l->first_free_entry;
            l->first_free_entry->prev = NULL;
            l->first_free_entry = entry->next;
        } else {
            entry = arena_push(&l->arena, struct entity_lookup_entry);
        }
        MEMZERO_STRUCT(entry);

        entry->key = key;
        entry->entity = handle;
        if (prev) {
            entry->prev = prev;
            prev->next = entry;
        }

        bucket->last = entry;
        *slot = entry;
    }
}

void entity_lookup_remove(struct entity_lookup *l, struct entity_lookup_entry *entry)
{
    struct entity_lookup_bucket *bucket = &l->buckets[entry->key.hash % l->num_buckets];
    struct entity_lookup_entry *prev = entry->prev;
    struct entity_lookup_entry *next = entry->next;

    if (prev) {
        prev->next = next;
    } else {
        bucket->first = next;
    }

    if (next) {
        next->prev = prev;
    } else {
        bucket->last = prev;
    }

    if (l->first_free_entry) {
        l->first_free_entry->prev = entry;
    }
    entry->next = l->first_free_entry;
    entry->prev = NULL;
    l->first_free_entry = entry;
}

struct entity_lookup_key entity_lookup_key_from_two_handles(struct entity_handle h0, struct entity_handle h1)
{
    struct entity_lookup_key key = ZI;
    struct string b0 = STRING_FROM_STRUCT(&h0);
    struct string b1 = STRING_FROM_STRUCT(&h1);
    key.hash = hash_fnv64(HASH_FNV64_BASIS, b0);
    key.hash = hash_fnv64(key.hash, b1);
    return key;
}

/* ========================== *
 * Activate
 * ========================== */

void entity_activate(struct entity *ent, u64 current_tick_id)
{
    entity_enable_prop(ent, ENTITY_PROP_ACTIVE);
    ent->activation_tick = current_tick_id;
    ++ent->continuity_gen;
}