#include "asset_cache.h"
#include "sys.h"
#include "string.h"
#include "memory.h"
#include "arena.h"
#include "scratch.h"
#include "util.h"
#include "work.h"
#include "log.h"

/* ========================== *
 * Global state
 * ========================== */

#define MAX_ASSETS 1024
#define ASSET_LOOKUP_TABLE_CAPACITY (MAX_ASSETS * 4)

GLOBAL struct {
    struct sys_rw_mutex lookup_rw_mutex;
    struct asset lookup[ASSET_LOOKUP_TABLE_CAPACITY];
    u64 num_assets;

    struct sys_rw_mutex store_rw_mutex;
    struct arena store_arena;

#if RTC
    /* Array of len `num_assets` pointing into populated entries of `lookup`. */
    struct asset *dbg_table[ASSET_LOOKUP_TABLE_CAPACITY];
    u64 dbg_table_count;
    struct sys_mutex dbg_table_mutex;
#endif
} G = { 0 }, DEBUG_ALIAS(G, G_asset_cache);

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

struct asset_cache_startup_receipt asset_cache_startup(struct work_startup_receipt *work_sr)
{
    (UNUSED)work_sr;

    /* Init lookup */
    G.lookup_rw_mutex = sys_rw_mutex_alloc();
    /* Init store */
    G.store_rw_mutex = sys_rw_mutex_alloc();
    G.store_arena = arena_alloc(GIGABYTE(64));
#if RTC
    /* Init debug */
    G.dbg_table_mutex = sys_mutex_alloc();
#endif

    return (struct asset_cache_startup_receipt) { 0 };
}

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

INTERNAL void refresh_dbg_table(void)
{
#if RTC
    sys_mutex_lock(&G.dbg_table_mutex);
    {
        MEMZERO_ARRAY(G.dbg_table);
        G.dbg_table_count = 0;
        for (u64 i = 0; i < ARRAY_COUNT(G.lookup); ++i) {
            struct asset *asset = &G.lookup[i];
            if (asset->hash != 0) {
                G.dbg_table[G.dbg_table_count++] = asset;
            }
        }
    }
    sys_mutex_unlock(&G.dbg_table_mutex);
#endif
}

/* Returns first matching slot or first empty slot if not found.
 * Check returned slot->hash != 0 for presence. */
INTERNAL struct asset *asset_cache_get_slot_assume_locked(struct string key, u64 hash)
{
    u64 index = hash % ARRAY_COUNT(G.lookup);
    while (true) {
        struct asset *slot = &G.lookup[index];
        if (slot->hash) {
            /* Occupied */
            if (hash == slot->hash && string_eq(key, slot->key)) {
                /* Matched slot */
                return slot;
            } else {
                ++index;
                if (index >= ARRAY_COUNT(G.lookup)) {
                    index = 0;
                }
            }
        } else {
            /* Empty slot */
            return slot;
        }
    }
}

u64 asset_cache_hash(struct string key)
{
    /* TODO: Better hash */
    return hash_fnv64(HASH_FNV64_BASIS, BUFFER_FROM_STRING(key));
}

/* `key` text is copied by this function
 *
 * Returns existing asset entry or inserts a new one.
 *
 * If is_first_touch (out parameter) is set to true, then the caller has
 * inserted the asset into the cache.
 *
 * */
struct asset *asset_cache_touch(struct string key, u64 hash, b32 *is_first_touch)
{
    struct asset *asset = NULL;
    if (is_first_touch) {
        *is_first_touch = false;
    }

    /* Lookup */
    {
        sys_rw_mutex_lock_shared(&G.lookup_rw_mutex);
        asset = asset_cache_get_slot_assume_locked(key, hash);
        sys_rw_mutex_unlock_shared(&G.lookup_rw_mutex);
    }

    /* Insert if not found */
    if (!asset->hash) {
        sys_rw_mutex_lock_exclusive(&G.lookup_rw_mutex);

        /* Re-check asset presence in case it was inserted since lock */
        asset = asset_cache_get_slot_assume_locked(key, hash);

        if (!asset->hash) {
            if (G.num_assets >= MAX_ASSETS) {
                sys_panic(STR("Max assets reached"));
            }
            struct string key_stored = { 0 };
            {
                /* Copy key to store */
                struct asset_cache_store store = asset_cache_store_open();
                key_stored = string_copy(store.arena, key);
                asset_cache_store_close(&store);
            }
            /* Initialize asset data */
            logf_info("Inserting asset cache entry for \"%F\"", FMT_STR(key));
            *asset = (struct asset) {
                .status = ASSET_STATUS_UNINITIALIZED,
                .hash = hash,
                .key = key_stored,
                .work_ready_sf = sync_flag_alloc(),
                .asset_ready_sf = sync_flag_alloc()
            };
            if (is_first_touch) {
                *is_first_touch = true;
            }
            ++G.num_assets;

            refresh_dbg_table();
        }

        sys_rw_mutex_unlock_exclusive(&G.lookup_rw_mutex);
    }

    return asset;
}

/* ========================== *
 * Marking
 * ========================== */

/* Call this once asset work has been created */
void asset_cache_mark_loading(struct asset *asset)
{
    asset->status = ASSET_STATUS_LOADING;
}

/* Call this once asset work has finished */
void asset_cache_mark_ready(struct asset *asset, void *store_data)
{
    asset->store_data = store_data;
    asset->status = ASSET_STATUS_READY;
    WRITE_BARRIER();
    sync_flag_set(&asset->asset_ready_sf);
}

/* ========================== *
 * Work
 * ========================== */

/* NOTE: If an asset doesn't have any load work then call this function with `NULL` */
void asset_cache_set_work(struct asset *asset, struct work_handle *handle)
{
    asset->work = handle ? *handle : (struct work_handle) { 0 };
    sync_flag_set(&asset->work_ready_sf);
}

void asset_cache_wait(struct asset *asset)
{
    if (asset->status != ASSET_STATUS_READY) {
        /* Wait for work to be set */
        sync_flag_wait(&asset->work_ready_sf);
        /* Help with work */
        if (asset->work.gen != 0) {
            work_help(asset->work);
        }
        /* Wait for asset to be ready */
        sync_flag_wait(&asset->asset_ready_sf);
    }
}

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

/* NOTE: At the moment only one global asset store exists, however in the
 * future there could be more based on asset lifetime. */

void *asset_cache_get_store_data(struct asset *asset)
{
    if (asset->status == ASSET_STATUS_READY) {
        return asset->store_data;
    } else {
        return NULL;
    }
}

/* Asset store should be opened to allocate memory to the store arena */
struct asset_cache_store asset_cache_store_open(void)
{
    struct asset_cache_store store = {
        .rw_mutex = &G.store_rw_mutex,
        .arena = &G.store_arena
    };
    sys_rw_mutex_lock_exclusive(store.rw_mutex);
    return store;
}

void asset_cache_store_close(struct asset_cache_store *store)
{
    sys_rw_mutex_unlock_exclusive(store->rw_mutex);
}