#include "ds.h"
#include "mem_arena.h"

#include <string.h>

#define RT_EMPTY_KEY 0ull
#define RT_TOMBSTONE 1ull

static uint64_t FixKey(uint64_t key) {
    if (key < RT_TOMBSTONE + 1)
        key = UINT64_MAX - key;
    return key;
}

RT_DLLEXPORT rt_hashtable rtCreateHashtable(uint64_t capacity,
                                            void *memory,
                                            rt_hashtable_grow_memory_fn *grow_memory_cb,
                                            void *userdata) {
    RT_ASSERT(RT_IS_POWER_OF_TWO(capacity), "Hashtable capacity must be a power of two.");
    memset(memory, 0, sizeof(uint64_t) * capacity);
    rt_hashtable hashtable = {
        .keys                 = memory,
        .values               = (uint64_t *)memory + capacity,
        .capacity             = capacity,
        .used_slots           = 0,
        .GrowMemoryCallback   = grow_memory_cb,
        .grow_memory_userdata = userdata,
    };
    return hashtable;
}

static rt_result Rehash(rt_hashtable *hashtable) {
    void *newmem = hashtable->GrowMemoryCallback(hashtable->keys,
                                                 2 * sizeof(uint64_t) * hashtable->capacity,
                                                 4 * sizeof(uint64_t) * hashtable->capacity,
                                                 hashtable->grow_memory_userdata);
    if (!newmem)
        return RT_OUT_OF_MEMORY;
    uint64_t old_capacity = hashtable->capacity;
    hashtable->capacity *= 2;
    hashtable->keys   = newmem;
    hashtable->values = hashtable->keys + hashtable->capacity;

    memset(hashtable->keys + old_capacity, 0, sizeof(uint64_t) * old_capacity);

    uint64_t *keys    = hashtable->keys;
    uint64_t capacity = hashtable->capacity;
    uint64_t mod      = capacity - 1;
    for (uint64_t i = 0; i < old_capacity; ++i) {
        if (keys[i] > RT_TOMBSTONE) {
            uint64_t key    = keys[i];
            uint64_t offset = 0;
            while (offset < capacity) {
                uint64_t new_slot = (key + offset) & mod;
                if (keys[new_slot] == key || keys[new_slot] <= RT_TOMBSTONE) {
                    if (new_slot != i) {
                        keys[new_slot]              = key;
                        hashtable->values[new_slot] = hashtable->values[i];
                        keys[i]                     = RT_TOMBSTONE;
                    }
                    break;
                }
            }
            if (offset == capacity) {
                rtReportError("CORE", "Rehashing failed. This is a bug.");
                return RT_UNKNOWN_ERROR;
            }
        }
    }
    return RT_SUCCESS;
}

RT_DLLEXPORT rt_result rtHashtableInsert(rt_hashtable *hashtable, uint64_t key, uint64_t value) {
    key = FixKey(key);

    if (hashtable->GrowMemoryCallback && (hashtable->used_slots > (hashtable->capacity >> 1))) {
        rt_result res = Rehash(hashtable);
        if (res != RT_SUCCESS)
            return res;
    }

    uint64_t *keys    = hashtable->keys;
    uint64_t capacity = hashtable->capacity;
    uint64_t mod      = capacity - 1;
    uint64_t offset   = 0;
    while (offset < capacity) {
        uint64_t index = (key + offset) & mod;
        ++offset;

        if (keys[index] == key || keys[index] <= RT_TOMBSTONE) {
            keys[index]              = key;
            hashtable->values[index] = value;
            ++hashtable->used_slots;
            return RT_SUCCESS;
        }
    }
    return RT_OUT_OF_MEMORY;
}

RT_DLLEXPORT uint64_t rtHashtableLookup(const rt_hashtable *hashtable,
                                        uint64_t key,
                                        uint64_t default_value) {
    key = FixKey(key);

    uint64_t capacity = hashtable->capacity;
    uint64_t mod      = capacity - 1;
    uint64_t offset   = 0;
    while (offset < capacity) {
        uint64_t index = (key + offset) & mod;
        ++offset;

        if (hashtable->keys[index] == key) {
            return hashtable->values[index];
        } else if (hashtable->keys[index] == RT_EMPTY_KEY) {
            break;
        }
    }
    return default_value;
}

RT_DLLEXPORT void rtHashtableRemove(rt_hashtable *hashtable, uint64_t key) {
    key = FixKey(key);
    uint64_t *keys    = hashtable->keys;
    uint64_t capacity = hashtable->capacity;
    uint64_t mod      = capacity - 1;
    uint64_t offset   = 0;
    while (offset < capacity) {
        uint64_t index = (key + offset) & mod;
        ++offset;

        if (keys[index] == key) {
            keys[index] = RT_TOMBSTONE;
        } else if (keys[index] == RT_EMPTY_KEY) {
            break;
        }
    }
}

RT_DLLEXPORT void *rtHashtableGrowMemoryFromArena(void *current,
                                                  size_t current_size,
                                                  size_t new_size,
                                                  void *userdata) {
    rt_arena *arena          = userdata;
    ptrdiff_t current_offset = (uint8_t *)current - (uint8_t *)arena->base;
    if (current_offset + current_size == arena->at) {
        /* The current allocation is the one at the end of the arena, so we can simply grow and be
         * done */
        if (!rtArenaPush(arena, new_size - current_size))
            return NULL;
        return current;
    } else {
        void *newmem = rtArenaPush(arena, new_size);
        if (newmem)
            memcpy(newmem, current, current_size);
        return newmem;
    }
}
