#include "../include/hashmap.h"
#include "../include/alloc.h"
#include <string.h>
#include <assert.h>

/* Optimized div by sizeof(mgui_hashmap_entry*) */
#if INTPTR_MAX == INT64_MAX
#define CAP_NUM_ENTRIES(cap) ((cap) >> 3)
#elif INTPTR_MAX == INT32_MAX
#define CAP_NUM_ENTRIES(cap) ((cap) >> 2)
#else
#error Unsupported system. Only systems where a pointer is 32-bits or 64-bits are supported.
#endif

#define HASH_TO_INDEX(hash) (hash & (CAP_NUM_ENTRIES(self->capacity)-1))
#define HASHMAP_ENTRY_GET_KEY(entry) ((char*)(entry) + sizeof(mgui_hashmap_entry))

/* |align| should be a multiple of 2 */
static size_t align_to(size_t value, size_t align) {
    const size_t is_aligned = (value & (align - 1));
    if(is_aligned == 0) {
        return value;
    } else {
        return (value & ~(align - 1)) + align;
    }
}

void mgui_hashmap_init(mgui_hashmap *self) {
    self->entries = NULL;
    self->capacity = 0;
    self->size = 0;
}

void mgui_hashmap_deinit(mgui_hashmap *self) {
    if(self->entries) {
        for(size_t i = 0; i < CAP_NUM_ENTRIES(self->capacity); ++i) {
            mgui_hashmap_entry *entry = self->entries[i];
            while(entry) {
                mgui_hashmap_entry *next = entry->next;
                mgui_free(entry);
                entry = next;
            }
        }

        mgui_free(self->entries);
        self->entries = NULL;
    }
    self->capacity = 0;
    self->size = 0;
}

static void mgui_hashmap_move_entry_to_front(mgui_hashmap *self, size_t index, mgui_hashmap_entry *new_entry) {
    new_entry->next = self->entries[index];
    self->entries[index] = new_entry;
}

static void mgui_hashmap_reorder_nodes(mgui_hashmap *self, size_t prev_capacity) {
    for(size_t i = 0; i < CAP_NUM_ENTRIES(prev_capacity); ++i) {
        mgui_hashmap_entry *entry = self->entries[i];
        mgui_hashmap_entry *prev_entry = NULL;
        while(entry) {
            const uint64_t new_index = HASH_TO_INDEX(entry->hash);
            mgui_hashmap_entry *next_entry = entry->next;

            if(new_index != i) {
                /* Remove entry by replacing this entry with the next entry */
                if(prev_entry)
                    prev_entry->next = next_entry;
                else
                    self->entries[i] = next_entry;

                /* Insert the entry at the new index */
                mgui_hashmap_move_entry_to_front(self, new_index, entry);
            } else {
                prev_entry = entry;
            }

            entry = next_entry;
        }
    }
}

static void mgui_hashmap_ensure_capacity(mgui_hashmap *self, size_t new_capacity) {
    if(self->capacity > new_capacity)
        return;

    size_t cap = self->capacity;
    if(cap == 0)
        cap = 8;

    while(cap < new_capacity) {
        cap = cap + (cap >> 1); /* *= 1.5 */
    }

    size_t prev_capacity = self->capacity;
    self->capacity = align_to(cap, sizeof(mgui_hashmap_entry*));
    self->entries = mgui_realloc(self->entries, self->capacity);

    memset((char*)self->entries + prev_capacity, 0, self->capacity - prev_capacity);
    mgui_hashmap_reorder_nodes(self, prev_capacity);
}

void mgui_hashmap_insert(mgui_hashmap *self, const char *key, size_t key_size, void *value, uint64_t *hash_out) {
    mgui_hashmap_ensure_capacity(self, (self->size + 1) * sizeof(mgui_hashmap_entry*));
    ++self->size;

    const uint64_t hash = mgui_hashmap_hash(key, key_size);
    const size_t index = HASH_TO_INDEX(hash);

#ifndef NDEBUG
    void *value_out;
    assert(!mgui_hashmap_get_by_hash(self, key, key_size, hash, &value_out));
#endif

    if(hash_out)
        *hash_out = hash;
    
    mgui_hashmap_entry *new_entry = mgui_alloc(sizeof(mgui_hashmap_entry) + key_size);
    new_entry->hash = hash;
    new_entry->value = value;
    new_entry->key_size = key_size;
    memcpy(HASHMAP_ENTRY_GET_KEY(new_entry), key, key_size);
    mgui_hashmap_move_entry_to_front(self, index, new_entry);
}

bool mgui_hashmap_get(mgui_hashmap *self, const char *key, size_t key_size, void **value_out) {
    return mgui_hashmap_get_by_hash(self, key, key_size, mgui_hashmap_hash(key, key_size), value_out);
}

bool mgui_hashmap_get_by_hash(mgui_hashmap *self, const char *key, size_t key_size, uint64_t hash, void **value_out) {
    assert(hash == mgui_hashmap_hash(key, key_size));
    if(!self->entries)
        return false;

    const size_t index = HASH_TO_INDEX(hash);
    mgui_hashmap_entry *entry = self->entries[index];
    while(entry) {
        if(hash == entry->hash && key_size == entry->key_size && memcmp(key, HASHMAP_ENTRY_GET_KEY(entry), key_size) == 0) {
            *value_out = entry->value;
            return true;
        }
        entry = entry->next;
    }
    return false;
}

/* TODO: Optimize by using xxHash instead */
/* fnv-1 hash */
uint64_t mgui_hashmap_hash(const char *data, size_t size) {
    const unsigned char *unsigned_data = (const unsigned char*)data;
    uint64_t hash = 14695981039346656037ULL;
    for(size_t i = 0; i < size; ++i)  {
        hash *= 1099511628211ULL;
        hash ^= unsigned_data[i];
    }
    return hash;
}

void mgui_hashmap_for_each(mgui_hashmap *self, bool(*callback)(void *value, void *userdata), void *userdata) {
    if(!self->entries)
        return;

    for(size_t i = 0; i < CAP_NUM_ENTRIES(self->capacity); ++i) {
        mgui_hashmap_entry *entry = self->entries[i];
        while(entry) {
            if(!callback(entry->value, userdata))
                return;
            entry = entry->next;
        }
    }
}

void mgui_hashmap_for_each_erase(mgui_hashmap *self, bool(*callback)(void *value, void *userdata), void *userdata) {
    if(!self->entries)
        return;

    for(size_t i = 0; i < CAP_NUM_ENTRIES(self->capacity); ++i) {
        mgui_hashmap_entry *entry = self->entries[i];
        mgui_hashmap_entry *prev_entry = NULL;
        while(entry) {
            mgui_hashmap_entry *next = entry->next;
            if(callback(entry->value, userdata)) {
                /* Remove entry by replacing this entry with the next entry */
                if(prev_entry)
                    prev_entry->next = next;
                else
                    self->entries[i] = next;

                mgui_free(entry);
            } else {
                prev_entry = entry;
            }
            entry = next;
        }
    }
}