#include "sim_snapshot.h"
#include "sim.h"
#include "sim_ent.h"
#include "sim_client.h"
#include "arena.h"
#include "scratch.h"

#define TICK_LOOKUP_BUCKETS 127
#define CLIENT_LOOKUP_BUCKETS 127

struct sim_snapshot_lookup_bucket {
    struct sim_snapshot *first;
    struct sim_snapshot *last;
};

/* ========================== *
 * Startup
 * ========================== */

GLOBAL struct {
    struct arena nil_arena;
    struct sim_snapshot_store *nil_snapshot_store;
    struct sim_snapshot *nil_snapshot;
    struct sim_ent *nil_ent;
    struct sim_client *nil_client;
} G = ZI, DEBUG_ALIAS(G, G_sim_snapshot);

/* Accessed via `sim_snapshot_store_nil()` */
READONLY struct sim_snapshot_store **_g_sim_snapshot_store_nil = &G.nil_snapshot_store;

/* Accessed via `sim_snapshot_nil()` */
READONLY struct sim_snapshot **_g_sim_snapshot_nil = &G.nil_snapshot;

/* Accessed via `sim_ent_nil()` */
READONLY struct sim_ent **_g_sim_ent_nil = &G.nil_ent;

/* Accessed via `sim_client_nil()` */
READONLY struct sim_client **_g_sim_client_nil = &G.nil_client;

struct sim_snapshot_startup_receipt sim_snapshot_startup(void)
{
    G.nil_arena = arena_alloc(GIGABYTE(1));

    /* Nil snapshot store */
    G.nil_snapshot_store = arena_push_zero(&G.nil_arena, struct sim_snapshot_store);
    G.nil_snapshot_store->valid = false;

    /* Nil snapshot */
    G.nil_snapshot = arena_push_zero(&G.nil_arena, struct sim_snapshot);
    G.nil_snapshot->valid = false;
    G.nil_snapshot->store = sim_snapshot_store_nil();

    /* Nil ent */
    G.nil_ent = arena_push_zero(&G.nil_arena, struct sim_ent);
    G.nil_ent->ss = sim_snapshot_nil();
    G.nil_ent->valid = false;
    G.nil_ent->handle = SIM_ENT_NIL_HANDLE;
    G.nil_ent->local_xform = XFORM_IDENT;
    G.nil_ent->cached_global_xform = XFORM_IDENT;
    G.nil_ent->cached_global_xform_dirty = false;
    G.nil_ent->friction = 0.5f;
    G.nil_ent->mass_unscaled = 1;
    G.nil_ent->inertia_unscaled = 1;
    G.nil_ent->sprite_local_xform = XFORM_IDENT;
    G.nil_ent->sprite_tint = COLOR_WHITE;

    /* Nil client */
    G.nil_client = arena_push_zero(&G.nil_arena, struct sim_client);
    G.nil_client->ss = sim_snapshot_nil();
    G.nil_client->valid = false;

    /* Lock nil arena */
    arena_set_readonly(&G.nil_arena);
    return (struct sim_snapshot_startup_receipt ) { 0 };
}

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

struct sim_snapshot_store *sim_snapshot_store_alloc(void)
{
    __prof;
    struct sim_snapshot_store *store;
    {
        struct arena arena = arena_alloc(GIGABYTE(64));
        store = arena_push_zero(&arena, struct sim_snapshot_store);
        store->arena = arena;
    }
    store->valid = true;
    store->num_lookup_buckets = TICK_LOOKUP_BUCKETS;
    store->lookup_buckets = arena_push_array_zero(&store->arena, struct sim_snapshot_lookup_bucket, store->num_lookup_buckets);
    return store;
}

void sim_snapshot_store_release(struct sim_snapshot_store *store)
{
    __prof;
    /* Release snapshot internal memory */
    for (u64 i = 0; i < store->num_lookup_buckets; ++i) {
        struct sim_snapshot_lookup_bucket *bucket = &store->lookup_buckets[i];
        struct sim_snapshot *ss = bucket->first;
        while (ss) {
            struct sim_snapshot *next = ss->next_in_bucket;
            arena_release(&ss->clients_arena);
            arena_release(&ss->ents_arena);
            arena_release(&ss->arena);
            ss = next;
        }
    }
    {
        struct sim_snapshot *ss = store->first_free_snapshot;
        while (ss) {
            struct sim_snapshot *next = ss->next_free;
            arena_release(&ss->clients_arena);
            arena_release(&ss->ents_arena);
            arena_release(&ss->arena);
            ss = next;
        }
    }
    /* Release store */
    arena_release(&store->arena);
}

/* Releases all snapshots with tick in range [start, end] */
void sim_snapshot_store_release_ticks_in_range(struct sim_snapshot_store *store, u64 start, u64 end)
{
    if (start > end) {
        u64 swp = start;
        start = end;
        end = swp;
    }

    struct sim_snapshot *ss = sim_snapshot_from_tick(store, store->first_tick);
    while (ss->valid) {
        u64 tick = ss->tick;
        u64 next_tick = ss->next_tick;
        if (tick >= start) {
            if (tick <= end) {
                sim_snapshot_release(ss);
            } else {
                break;
            }
        }
        ss = sim_snapshot_from_tick(store, next_tick);
    }
}

/* ========================== *
 * Alloc
 * ========================== */

/* Produces a new snapshot at `tick` with data copied from `src` snapshot. */
struct sim_snapshot *sim_snapshot_alloc(struct sim_snapshot_store *store, struct sim_snapshot *src, u64 tick)
{
    if (tick == 0) {
        return sim_snapshot_nil();
    }

    struct sim_snapshot *ss;
    {
        struct arena arena = ZI;
        struct arena clients_arena = ZI;
        struct arena ents_arena = ZI;
        {
            ss = store->first_free_snapshot;
            if (ss) {
                /* Re-use existing snasphot arenas */
                store->first_free_snapshot = ss->next_free;
                ents_arena = ss->ents_arena;
                clients_arena = ss->clients_arena;
                arena = ss->arena;
            } else {
                /* Arenas allocated here will be released along with the entire snasphot store */
                arena = arena_alloc(GIGABYTE(8));
                clients_arena = arena_alloc(GIGABYTE(8));
                ents_arena = arena_alloc(GIGABYTE(8));
            }
        }
        arena_reset(&arena);
        ss = arena_push_zero(&arena, struct sim_snapshot);
        ss->arena = arena;

        ss->clients_arena = clients_arena;
        arena_reset(&ss->clients_arena);

        ss->ents_arena = ents_arena;
        arena_reset(&ss->ents_arena);
    }

    ss->tick = tick;
    ss->valid = true;
    ss->store = store;
    ++store->num_ticks;

    /* Copy src info */
    ss->is_master = src->is_master;
    ss->real_dt_ns = src->real_dt_ns;
    ss->real_time_ns = src->real_time_ns;
    ss->world_timescale = src->world_timescale;
    ss->world_dt_ns = src->world_dt_ns;
    ss->world_time_ns = src->world_time_ns;
    ss->continuity_gen = src->continuity_gen;
    ss->local_client = src->local_client;

    /* Copy client lookup buckets */
    if (src->num_client_lookup_buckets > 0) {
        ss->num_client_lookup_buckets = src->num_client_lookup_buckets;
        ss->client_lookup_buckets = arena_push_array(&ss->arena, struct sim_client_lookup_bucket, ss->num_client_lookup_buckets);
        MEMCPY(ss->client_lookup_buckets, src->client_lookup_buckets, sizeof(*ss->client_lookup_buckets) * ss->num_client_lookup_buckets);
    } else {
        ss->num_client_lookup_buckets = CLIENT_LOOKUP_BUCKETS;
        ss->client_lookup_buckets = arena_push_array_zero(&ss->arena, struct sim_client_lookup_bucket, ss->num_client_lookup_buckets);
    }

    /* Copy clients */
    ss->first_free_client = src->first_free_client;
    ss->num_clients_allocated = src->num_clients_allocated;
    ss->num_clients_reserved = src->num_clients_reserved;
    ss->clients = arena_push_array(&ss->clients_arena, struct sim_client, ss->num_clients_reserved);
    for (u64 i = 0; i < ss->num_clients_reserved; ++i) {
        struct sim_client *dst_client = &ss->clients[i];
        struct sim_client *src_client = &src->clients[i];
        *dst_client = *src_client;
        dst_client->ss = ss;
    }

    /* Copy entities */
    ss->first_free_ent = src->first_free_ent;
    ss->num_ents_allocated = src->num_ents_allocated;
    ss->num_ents_reserved = src->num_ents_reserved;
    ss->ents = arena_push_array(&ss->ents_arena, struct sim_ent, ss->num_ents_reserved);
    if (ss->num_ents_reserved == 0) {
        /* Create root ent if copying from nil store */
        struct sim_ent *root = arena_push(&ss->ents_arena, struct sim_ent);
        *root = *sim_ent_nil();
        root->ss = ss;
        root->handle = SIM_ENT_ROOT_HANDLE;
        root->valid = true;
        root->is_root = true;
        ++ss->num_ents_allocated;
        ++ss->num_ents_reserved;
    } else {
        for (u64 i = 0; i < ss->num_ents_reserved; ++i) {
            struct sim_ent *dst_ent = &ss->ents[i];
            struct sim_ent *src_ent = &src->ents[i];
            *dst_ent = *src_ent;
            dst_ent->ss = ss;
        }
    }

    /* Release duplicate tick if it exists */
    {
        struct sim_snapshot *existing = sim_snapshot_from_tick(store, tick);
        if (existing->valid) {
            sim_snapshot_release(existing);
        }
    }

    /* Linear search to insert snapshot in tick order */
    {
        struct sim_snapshot *prev = sim_snapshot_from_tick(store, store->last_tick);
        while (prev->valid) {
            if (prev->tick < tick) {
                break;
            }
            prev = sim_snapshot_from_tick(store, prev->prev_tick);
        }
        if (prev->valid) {
            struct sim_snapshot *next = sim_snapshot_from_tick(store, prev->next_tick);
            if (next->valid) {
                next->prev_tick = tick;
            } else {
                store->last_tick = tick;
            }
            ss->next_tick = prev->next_tick;
            ss->prev_tick = prev->tick;
            prev->next_tick = ss->tick;
        } else {
            struct sim_snapshot *first = sim_snapshot_from_tick(store, store->first_tick);
            if (first->valid) {
                ss->next_tick = first->tick;
                first->prev_tick = tick;
            } else {
                store->last_tick = tick;
            }
            ss->next_tick = store->first_tick;
            store->first_tick = tick;
        }
    }

    /* Insert into lookup */
    {
        u64 bucket_index = tick % store->num_lookup_buckets;
        struct sim_snapshot_lookup_bucket *bucket = &store->lookup_buckets[bucket_index];
        if (bucket->last) {
            bucket->last->next_in_bucket = ss;
            ss->prev_in_bucket = bucket->last;
        } else {
            bucket->first = ss;
        }
        bucket->last = ss;
    }

    return ss;
}

void sim_snapshot_release(struct sim_snapshot *ss)
{
    struct sim_snapshot_store *store = ss->store;

    /* Remove from lookup */
    {
        u64 bucket_index = ss->tick % store->num_lookup_buckets;
        struct sim_snapshot_lookup_bucket *bucket = &store->lookup_buckets[bucket_index];
        struct sim_snapshot *prev = ss->prev_in_bucket;
        struct sim_snapshot *next = ss->next_in_bucket;
        if (prev) {
            prev->next_in_bucket = next;
        } else {
            bucket->first = next;
        }
        if (next) {
            next->prev_in_bucket = prev;
        } else {
            bucket->last = prev;
        }
    }

    /* Remove from snapshot list */
    {
        struct sim_snapshot *prev = sim_snapshot_from_tick(store, ss->prev_tick);
        struct sim_snapshot *next = sim_snapshot_from_tick(store, ss->next_tick);
        if (prev->valid) {
            prev->next_tick = next->tick;
        } else {
            store->first_tick = next->tick;
        }
        if (next->valid) {
            next->prev_tick = prev->tick;
        } else {
            store->last_tick = prev->tick;
        }
    }

    ss->valid = false;
    ss->next_free = store->first_free_snapshot;
    store->first_free_snapshot = ss;
    --store->num_ticks;
}

/* ========================== *
 * Lookup
 * ========================== */

struct sim_snapshot *sim_snapshot_from_tick(struct sim_snapshot_store *store, u64 tick)
{
    struct sim_snapshot *ss = sim_snapshot_nil();
    if (tick > 0) {
        u64 bucket_index = tick % store->num_lookup_buckets;
        struct sim_snapshot_lookup_bucket *bucket = &store->lookup_buckets[bucket_index];
        for (struct sim_snapshot *search = bucket->first; search; search = search->next_in_bucket) {
            if (search->tick == tick) {
                ss = search;
                break;
            }
        }
    }
    return ss;
}

/* ========================== *
 * Lerp
 * ========================== */

struct sim_snapshot *sim_snapshot_alloc_from_lerp(struct sim_snapshot_store *store, struct sim_snapshot *ss0, struct sim_snapshot *ss1, f64 blend)
{
    __prof;

    /* New snapshot will be allocated with same tick as ss0 or ss1, so the result should go into a different store */
    ASSERT(ss0->store != store && ss1->store != store);

    struct sim_snapshot *ss;
    b32 should_blend = true;
    if (ss0->continuity_gen == ss1->continuity_gen && 0 < blend && blend < 1) {
        ss = sim_snapshot_alloc(store, ss0, ss0->tick);
    } else if (math_round_to_int64(blend) <= 0) {
        ss = sim_snapshot_alloc(store, ss0, ss0->tick);
        should_blend = false;
    } else {
        ss = sim_snapshot_alloc(store, ss1, ss1->tick);
        should_blend = false;
    }

    if (!ss0->valid || !ss1->valid) {
        /* New snapshot allocation caused one of the src snapshots original to release.
         * ss0 & ss1 should be from a separate store than the allocating one. */
        ASSERT(false);
    }

    if (should_blend) {
        /* Blend time */
        ss->real_dt_ns = math_lerp_i64(ss0->real_dt_ns, ss1->real_dt_ns, blend);
        ss->real_time_ns = math_lerp_i64(ss0->real_time_ns, ss1->real_time_ns, blend);
        ss->world_timescale = math_lerp_f64(ss0->world_timescale, ss1->world_timescale, blend);
        ss->world_dt_ns = math_lerp_i64(ss0->world_dt_ns, ss1->world_dt_ns, blend);
        ss->world_time_ns = math_lerp_i64(ss0->world_time_ns, ss1->world_time_ns, blend);

        /* Blend clients */
        {
            __profscope(snapshot_lerp_clients);
            u64 num_clients = min_u64(ss0->num_clients_reserved, ss1->num_clients_reserved);
            for (u64 i = 0; i < num_clients; ++i) {
                struct sim_client *c = &ss->clients[i];
                struct sim_client *c0 = &ss0->clients[i];
                struct sim_client *c1 = &ss1->clients[i];
                sim_client_lerp(c, c0, c1, blend);
            }
        }

        /* Blend entities */
        {
            __profscope(snapshot_lerp_entities);
            u64 num_entities = min_u64(ss0->num_ents_reserved, ss1->num_ents_reserved);
            for (u64 i = 0; i < num_entities; ++i) {
                struct sim_ent *e = &ss->ents[i];
                struct sim_ent *e0 = &ss0->ents[i];
                struct sim_ent *e1 = &ss1->ents[i];
                sim_ent_lerp(e, e0, e1, blend);
            }
        }
    }


    return ss;
}

/* ========================== *
 * Encode
 * ========================== */

void sim_snapshot_encode(struct bitbuff_writer *bw, struct sim_client *receiver, struct sim_snapshot *ss0, struct sim_snapshot *ss1)
{
    __prof;

    /* TODO: Don't encode this */
    bw_write_bit(bw, ss1->is_master);

    bw_write_iv(bw, ss1->real_dt_ns);
    bw_write_iv(bw, ss1->real_time_ns);

    bw_write_f64(bw, ss1->world_timescale);
    bw_write_iv(bw, ss1->world_dt_ns);
    bw_write_iv(bw, ss1->world_time_ns);

    bw_write_uv(bw, ss1->continuity_gen);
    bw_write_uv(bw, ss1->phys_iteration);

    bw_write_uv(bw, receiver->handle.gen);
    bw_write_uv(bw, receiver->handle.idx);

    /* Clients */
    if (ss1->num_clients_allocated == ss0->num_clients_allocated) {
        bw_write_bit(bw, 0);
    } else {
        bw_write_bit(bw, 1);
        bw_write_uv(bw, ss1->num_clients_allocated);
    }
    if (ss1->num_clients_reserved == ss0->num_clients_reserved) {
        bw_write_bit(bw, 0);
    } else {
        bw_write_bit(bw, 1);
        bw_write_uv(bw, ss1->num_clients_reserved);
    }
    for (u64 i = 0; i < ss1->num_clients_reserved; ++i) {
        struct sim_client *c0 = sim_client_nil();
        if (i < ss0->num_clients_reserved) {
            c0 = &ss0->clients[i];
        }
        struct sim_client *c1 = &ss1->clients[i];
        sim_client_encode(bw, c0, c1);
    }

    /* Ents */
    if (ss1->num_ents_allocated == ss0->num_ents_allocated) {
        bw_write_bit(bw, 0);
    } else {
        bw_write_bit(bw, 1);
        bw_write_uv(bw, ss1->num_ents_allocated);
    }
    if (ss1->num_ents_reserved == ss0->num_ents_reserved) {
        bw_write_bit(bw, 0);
    } else {
        bw_write_bit(bw, 1);
        bw_write_uv(bw, ss1->num_ents_reserved);
    }
    for (u64 i = 0; i < ss1->num_ents_reserved; ++i) {
        struct sim_ent *e0 = sim_ent_nil();
        if (i < ss0->num_ents_reserved) {
            e0 = &ss0->ents[i];
        }
        struct sim_ent *e1 = &ss1->ents[i];
        sim_ent_encode(bw, e0, e1);
    }
}

/* ========================== *
 * Decode
 * ========================== */

void sim_snapshot_decode(struct bitbuff_reader *br, struct sim_snapshot *ss)
{
    __prof;

    /* TODO: Don't encode this */
    ss->is_master = br_read_bit(br);

    ss->real_dt_ns = br_read_iv(br);
    ss->real_time_ns = br_read_iv(br);

    ss->world_timescale = br_read_f64(br);
    ss->world_dt_ns = br_read_iv(br);
    ss->world_time_ns = br_read_iv(br);

    ss->continuity_gen = br_read_uv(br);
    ss->phys_iteration = br_read_uv(br);

    ss->local_client.gen = br_read_uv(br);
    ss->local_client.idx = br_read_uv(br);

    /* Clients */
    if (br_read_bit(br)) {
        ss->num_clients_allocated = br_read_uv(br);
    }
    if (br_read_bit(br)) {
        u64 old_num_clients_reserved = ss->num_clients_reserved;
        ss->num_clients_reserved = br_read_uv(br);
        i64 reserve_diff = (i64)ss->num_clients_reserved - (i64)old_num_clients_reserved;
        if (reserve_diff > 0) {
            arena_push_array(&ss->clients_arena, struct sim_client, reserve_diff);
            for (u64 i = old_num_clients_reserved; i < ss->num_clients_reserved; ++i) {
                struct sim_client *c = &ss->clients[i];
                *c = *sim_client_nil();
                c->ss = ss;
            }
        }
    }
    for (u64 i = 0; i < ss->num_clients_reserved; ++i) {
        struct sim_client *c = &ss->clients[i];
        sim_client_decode(br, c);
    }

    /* Ents */
    if (br_read_bit(br)) {
        ss->num_ents_allocated = br_read_uv(br);
    }
    if (br_read_bit(br)) {
        u64 old_num_ents_reserved = ss->num_ents_reserved;
        ss->num_ents_reserved = br_read_uv(br);
        i64 reserve_diff = (i64)ss->num_ents_reserved - (i64)old_num_ents_reserved;
        if (reserve_diff > 0) {
            arena_push_array(&ss->ents_arena, struct sim_ent, reserve_diff);
            for (u64 i = old_num_ents_reserved; i < ss->num_ents_reserved; ++i) {
                struct sim_ent *e = &ss->ents[i];
                *e = *sim_ent_nil();
                e->ss = ss;
            }
        }
    }
    for (u64 i = 0; i < ss->num_ents_reserved; ++i) {
        struct sim_ent *e = &ss->ents[i];
        sim_ent_decode(br, e);
    }
}