#include "sim_ent.h"
#include "sim.h"
#include "math.h"
#include "bitbuff.h"

/* ========================== *
 * Ent allocation
 * ========================== */

INTERNAL struct sim_ent *sim_ent_alloc_internal(struct sim_snapshot *ss)
{
    struct sim_ent *ent = NULL;
    struct sim_ent_handle handle = ZI;
    if (ss->first_free_ent.gen) {
        /* Reuse from free list */;
        ent = sim_ent_from_handle(ss, ss->first_free_ent);
        handle = ent->handle;
        ++handle.gen;
        ss->first_free_ent = ent->next_free;
    } else {
        /* Make new */
        ent = arena_push(&ss->ents_arena, struct sim_ent);
        handle = (struct sim_ent_handle) { .gen = 1, .idx = ss->num_ents_reserved++ };
    }
    *ent = *sim_ent_nil();
    ent->ss = ss;
    ent->valid = true;
    ent->handle = handle;
    ent->_is_xform_dirty = true;
    ++ss->num_ents_allocated;
    return ent;
}

struct sim_ent *sim_ent_alloc(struct sim_ent *parent)
{
    ASSERT(parent->valid);
    struct sim_snapshot *ss = parent->ss;
    struct sim_ent *e = sim_ent_alloc_internal(ss);
    sim_ent_link_parent(e, parent);
    return e;
}

INTERNAL void sim_ent_release_internal(struct sim_ent *ent)
{
    struct sim_snapshot *ss = ent->ss;
    /* Release children */
    struct sim_ent_handle first_handle = ent->first;
    if (first_handle.gen) {
        for (struct sim_ent *child = sim_ent_from_handle(ss, first_handle); child->valid; child = sim_ent_from_handle(ss, child->next)) {
            sim_ent_release_internal(child);
        }
    }

    /* Release */
    ++ent->handle.gen;
    ent->valid = false;
    ent->next_free = ss->first_free_ent;
    ss->first_free_ent = ent->handle;
    --ss->num_ents_allocated;
}

void sim_ent_release(struct sim_ent *ent)
{
    if (ent->parent.gen) {
        sim_ent_unlink_from_parent(ent);
    }
    sim_ent_release_internal(ent);
}

/* ========================== *
 * Ent query
 * ==========================  */

/* Returns a valid ent or read-only nil ent. Always safe to read result, need to check `valid` to write. */
struct sim_ent *sim_ent_from_handle(struct sim_snapshot *ss, struct sim_ent_handle handle)
{
    if (handle.gen != 0 && handle.idx < ss->num_ents_reserved) {
        struct sim_ent *ent = &ss->ents[handle.idx];
        if (ent->handle.gen == handle.gen) {
            return ent;
        }
    }
    return sim_ent_nil();
}

struct sim_ent *sim_ent_find_first_match_one(struct sim_snapshot *ss, enum sim_ent_prop prop)
{
    u64 count = ss->num_ents_reserved;
    struct sim_ent *entities = ss->ents;
    for (u64 ent_index = 0; ent_index < count; ++ent_index) {
        struct sim_ent *ent = &entities[ent_index];
        if (ent->valid && sim_ent_has_prop(ent, prop)) {
            return ent;
        }
    }
    return sim_ent_nil();
}

struct sim_ent *sim_ent_find_first_match_all(struct sim_snapshot *ss, struct sim_ent_prop_array props)
{
    u64 count = ss->num_ents_reserved;
    struct sim_ent *entities = ss->ents;
    for (u64 ent_index = 0; ent_index < count; ++ent_index) {
        struct sim_ent *ent = &entities[ent_index];
        if (ent->valid) {
            b32 all = true;
            for (u64 i = 0; i < props.count; ++i) {
                if (!sim_ent_has_prop(ent, props.props[i])) {
                    all = false;
                    break;
                }
            }
            if (all) {
                return ent;
            }
        }
    }
    return sim_ent_nil();
}

/* ========================== *
 * Ent tree
 * ========================== */

void sim_ent_link_parent(struct sim_ent *ent, struct sim_ent *parent)
{
    struct sim_snapshot *ss = ent->ss;

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

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

    ent->parent = parent_handle;

    struct sim_ent_handle last_child_handle = parent->last;
    struct sim_ent *last_child = sim_ent_from_handle(ss, 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 ent. */
void sim_ent_unlink_from_parent(struct sim_ent *ent)
{
    struct sim_snapshot *ss = ent->ss;

    struct sim_ent_handle parent_handle = ent->parent;
    struct sim_ent *parent = sim_ent_from_handle(ss, parent_handle);
    struct sim_ent *prev = sim_ent_from_handle(ss, ent->prev);
    struct sim_ent *next = sim_ent_from_handle(ss, 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 = sim_ent_nil()->handle;
    ent->next = sim_ent_nil()->handle;
}

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

void sim_ent_activate(struct sim_ent *ent, u64 current_tick)
{
    sim_ent_enable_prop(ent, SIM_ENT_PROP_ACTIVE);
    ent->activation_tick = current_tick;
    ++ent->continuity_gen;
}

/* ========================== *
 * Ent xform
 * ========================== */

INTERNAL void sim_ent_mark_child_xforms_dirty(struct sim_snapshot *ss, struct sim_ent *ent)
{
    for (struct sim_ent *child = sim_ent_from_handle(ss, ent->first); child->valid; child = sim_ent_from_handle(ss, child->next)) {
        if (child->_is_xform_dirty) {
            break;
        } else {
            child->_is_xform_dirty = true;
            sim_ent_mark_child_xforms_dirty(ss, child);
        }
    }
}

INTERNAL struct xform sim_ent_get_xform_internal(struct sim_snapshot *ss, struct sim_ent *ent)
{
    struct xform xf;
    if (ent->_is_xform_dirty) {
        if (ent->is_top) {
            xf = ent->_local_xform;
        } else {
            struct sim_ent *parent = sim_ent_from_handle(ss, ent->parent);
            xf = sim_ent_get_xform_internal(ss, parent);
            xf = xform_mul(xf, ent->_local_xform);
            ent->_xform = xf;
            ent->_is_xform_dirty = false;
        }
        ent->_xform = xf;
        ent->_is_xform_dirty = false;
    } else {
        xf = ent->_xform;
    }
    return xf;
}

struct xform sim_ent_get_xform(struct sim_ent *ent)
{
    struct xform xf;
    if (ent->_is_xform_dirty) {
        if (ent->is_top) {
            xf = ent->_local_xform;
        } else {
            struct sim_snapshot *ss = ent->ss;
            struct sim_ent *parent = sim_ent_from_handle(ss, ent->parent);
            xf = sim_ent_get_xform_internal(ss, parent);
            xf = xform_mul(xf, ent->_local_xform);
            ent->_xform = xf;
            ent->_is_xform_dirty = false;
        }
        ent->_xform = xf;
        ent->_is_xform_dirty = false;
    } else {
        xf = ent->_xform;
    }
    return xf;
}

struct xform sim_ent_get_local_xform(struct sim_ent *ent)
{
    return ent->_local_xform;
}

void sim_ent_set_xform(struct sim_ent *ent, struct xform xf)
{
    if (!xform_eq(xf, ent->_xform)) {
        struct sim_snapshot *ss = ent->ss;
        /* Update local xform */
        if (ent->is_top) {
            ent->_local_xform = xf;
        } else {
            struct sim_ent *parent = sim_ent_from_handle(ss, ent->parent);
            struct xform parent_global = sim_ent_get_xform_internal(ss, parent);
            ent->_local_xform = xform_mul(xform_invert(parent_global), xf);
        }
        ent->_xform = xf;
        ent->_is_xform_dirty = false;
        sim_ent_mark_child_xforms_dirty(ss, ent);
    }
}

void sim_ent_set_local_xform(struct sim_ent *ent, struct xform xf)
{
    if (!xform_eq(xf, ent->_local_xform)) {
        ent->_local_xform = xf;
        ent->_is_xform_dirty = true;
        sim_ent_mark_child_xforms_dirty(ent->ss, ent);
    }
}

/* ========================== *
 * Ent movement
 * ========================== */

void sim_ent_apply_linear_impulse(struct sim_ent *ent, struct v2 impulse, struct v2 point)
{
    struct xform xf = sim_ent_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);
    sim_ent_set_linear_velocity(ent, v2_add(ent->linear_velocity, v2_mul(impulse, inv_mass)));
    sim_ent_set_angular_velocity(ent, v2_wedge(vcp, impulse) * inv_inertia);
}

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

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

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

void sim_ent_apply_angular_impulse(struct sim_ent *ent, f32 impulse)
{
    struct xform xf = sim_ent_get_xform(ent);
    f32 scale = math_fabs(xform_get_determinant(xf));
    f32 inv_inertia = 1.f / (ent->inertia_unscaled * scale);
    sim_ent_set_angular_velocity(ent, ent->angular_velocity + impulse * inv_inertia);
}

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

/* ========================== *
 * Ent lerp
 * ========================== */

void sim_ent_lerp(struct sim_ent *e, struct sim_ent *e0, struct sim_ent *e1, f64 blend)
{
    if (sim_ent_is_valid_and_active(e0) && sim_ent_is_valid_and_active(e1)
        && e0->handle.gen == e1->handle.gen
        && e0->continuity_gen == e1->continuity_gen) {
        e->_local_xform = xform_lerp(e0->_local_xform, e1->_local_xform, blend);

        if (e->is_top) {
            /* TODO: Cache parent & child xforms in sim */
            struct xform e0_xf = sim_ent_get_xform(e0);
            struct xform e1_xf = sim_ent_get_xform(e1);
            sim_ent_set_xform(e, xform_lerp(e0_xf, e1_xf, blend));
        }

        e->control_force = math_lerp_f32(e0->control_force, e1->control_force, blend);
        e->control_torque = math_lerp_f32(e0->control_torque, e1->control_torque, blend);

        e->linear_velocity = v2_lerp(e0->linear_velocity, e1->linear_velocity, blend);
        e->angular_velocity = math_lerp_angle(e0->angular_velocity, e1->angular_velocity, blend);

        e->control.move = v2_lerp(e0->control.move, e1->control.move, blend);
        e->control.focus = v2_lerp(e0->control.focus, e1->control.focus, blend);

        e->sprite_local_xform = xform_lerp(e0->sprite_local_xform, e1->sprite_local_xform, blend);
        e->animation_last_frame_change_time_ns = math_lerp_i64(e0->animation_last_frame_change_time_ns, e1->animation_last_frame_change_time_ns, (f64)blend);
        e->animation_frame = (u32)math_round_to_int(math_lerp_f32(e0->animation_frame, e1->animation_frame, blend));

        e->camera_quad_xform = xform_lerp(e0->camera_quad_xform, e1->camera_quad_xform, blend);
        e->camera_xform_target = xform_lerp(e0->camera_xform_target, e1->camera_xform_target, blend);
        e->shake = math_lerp_f32(e0->shake, e1->shake, blend);

        e->tracer_gradient_start = v2_lerp(e0->tracer_gradient_start, e1->tracer_gradient_start, blend);
        e->tracer_gradient_end = v2_lerp(e0->tracer_gradient_end, e1->tracer_gradient_end, blend);
    }
}

/* ========================== *
 * Ent encode
 * ========================== */

void sim_ent_encode(struct bitbuff_writer *bw, struct sim_ent *e0, struct sim_ent *e1)
{
    struct sim_snapshot *ss = e1->ss;


    /* FIXME: Things like xforms need to be retreived manually rather than memcopied.
     *  This will also be true for things like ent handles once netids are implemented. */

    /* TODO: Granular delta encoding */

    u64 pos = 0;
    e1->ss = e0->ss;
    while (pos < sizeof(*e1)) {
        u64 chunk_size = min_u64(pos + 8, sizeof(*e1)) - pos;
        u8 *chunk0 = (u8 *)e0 + pos;
        u8 *chunk1 = (u8 *)e1 + pos;
        if (MEMEQ(chunk0, chunk1, chunk_size)) {
            bw_write_bit(bw, 0);
        } else {
            bw_write_bit(bw, 1);
            u64 bits = 0;
            MEMCPY(&bits, chunk1, chunk_size);
            bw_write_ubits(bw, bits, 64);
        }
        pos += 8;
    }
    e1->ss = ss;
}

/* ========================== *
 * Ent decode
 * ========================== */

void sim_ent_decode(struct bitbuff_reader *br, struct sim_ent *e)
{
    struct sim_snapshot *ss = e->ss;

    u64 pos = 0;
    while (pos < sizeof(*e)) {
        u8 *chunk = (u8 *)e + pos;
        if (br_read_bit(br)) {
            u64 chunk_size = min_u64(pos + 8, sizeof(*e)) - pos;
            u64 bits = br_read_ubits(br, 64);
            MEMCPY(chunk, &bits, chunk_size);
        }
        pos += 8;
    }

    e->ss = ss;
}