#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 = {
    0
};

/* 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
};

GLOBAL READONLY struct entity g_entity_default = {
    .valid = true,
    .local_xform = XFORM_IDENT_NOCAST,
    .cached_global_xform = XFORM_IDENT_NOCAST,
    .cached_global_xform_dirty = true,
    .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->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;
}

/* ========================== *
 * 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_default;
    entity->handle = handle;
    ++store->allocated;
    return entity;
}

struct entity *entity_alloc(struct entity *parent)
{
    ASSERT(parent->valid);
    struct entity_store *store = entity_get_store(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_parent(ent);
    }
    entity_release_internal(store, ent);
}

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

struct entity_store *entity_get_store(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_get_store(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_get_store(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_get_store(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);
    ent->linear_velocity = v2_add(ent->linear_velocity, v2_mul(impulse, inv_mass));
    ent->angular_velocity += v2_wedge(vcp, impulse) * inv_inertia;
}

#if 0
void entity_apply_force(struct entity *ent, struct v2 impulse, struct v2 world_point)
{
}
#endif

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);

    ent->linear_velocity = 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);

    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_get_store(ent);

    if (ent->parent.gen) {
        /* Unlink from current parent */
        entity_unlink_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;

    ent->is_top = parent->is_root;
}

/* NOTE: Entity will be dangling after calling this, should re-link to root entity. */
void entity_unlink_parent(struct entity *ent)
{
    struct entity_store *store = entity_get_store(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_nil_handle();
    ent->next = entity_nil_handle();
}