#include "map.h"
#include "crash.h"
#include <stdlib.h>
#include <assert.h>

#define DEFAULT_CAPACITY 16

size_t hash_bytes(const void* start, size_t size) {
    size_t hash = 0;
    const char* raw_data = start;
    for (size_t i = 0; i < size; i++) {
        hash = (hash << 5) - hash + raw_data[i];
    }
    return hash;
}

void map_init(
    map_t* map,
    hash_func_t hash_func,
    eq_func_t eq_func,
    double load_limit
) {
    map_init_capacity(
        map, hash_func, eq_func, load_limit, DEFAULT_CAPACITY);
}

void map_init_capacity(
    map_t* map,
    hash_func_t hash_func,
    eq_func_t eq_func,
    double load_limit,
    size_t capacity
) {
    map->hash_func = hash_func;
    map->eq_func = eq_func;
    map->load_limit = load_limit;
    map->size = 0;
    map->__num_buckets = capacity / load_limit + 1;
    map->__buckets = calloc(map->__num_buckets, sizeof(struct __map_entry));

    if (map->__buckets == NULL)
        crash("Out of memory allocating %ld hash buckets\n");
}

void map_destroy(map_t* map) {
    for (size_t i = 0; i < map->__num_buckets; i++) {
        struct __map_entry* node = map->__buckets[i].next;
        while (node != NULL) {
            struct __map_entry* next = node->next;
            free(node);
            node = next;
        }
    }

    free(map->__buckets);
}

/* impl detail: next is never null if the value is populated (put will alloc) */
static struct __map_entry* fetch_entry(const map_t* map, const void* key) {
    struct __map_entry* entry =
        &map->__buckets[map->hash_func(key) % map->__num_buckets];

    while (entry->next != NULL && !map->eq_func(key, entry->key)) {
        entry = entry->next;
    }
    return entry;
}

bool map_contains(const map_t* map, const void *key) {
    return fetch_entry(map, key)->next != NULL;
}

void* map_get(const map_t* map, const void* key) {
    return map_get_or_default(map, key, NULL);
}

void* map_get_or_default(const map_t* map, const void* key, void* default_val) {
    struct __map_entry* entry = fetch_entry(map, key);
    return entry->next == NULL ? default_val : entry->value;
}

static void insert_entry(map_t* map, struct __map_entry* entry) {
    struct __map_entry* slot = fetch_entry(map, entry->key);
    assert(slot->next == NULL);
    slot->key = entry->key;
    slot->value = entry->value;
    slot->next = entry;
    entry->next = NULL;
}

static void rehash(map_t* map) {
    size_t old_num_buckets = map->__num_buckets;
    struct __map_entry* old_buckets = map->__buckets;
    map->__num_buckets <<= 1;
    map->__buckets = calloc(map->__num_buckets, sizeof(struct __map_entry));

    for (size_t i = 0; i < old_num_buckets; i++) {
        struct __map_entry* old_entry = &old_buckets[i];
        if (old_entry->next == NULL) continue;

        struct __map_entry* cur = old_entry->next;
        while (cur->next != NULL) {
            struct __map_entry* next = cur->next;
            insert_entry(map, cur);
            cur = next;
        }

        struct __map_entry* new_entry = fetch_entry(map, old_entry->key);
        new_entry->key = old_entry->key;
        new_entry->value = old_entry->value;
        new_entry->next = cur;
    }

    free(old_buckets);
}

static void insert_allocating(
    map_t* map,
    struct __map_entry* entry,
    void* key,
    void* val
) {
    if ((double) ++map->size / map->__num_buckets > map->load_limit) {
        rehash(map);
        entry = fetch_entry(map, key);
    }

    struct __map_entry* next = calloc(1, sizeof(struct __map_entry));
    if (next == NULL)
        crash("Out of memory allocating hash entry\n");

    entry->key = key;
    entry->value = val;
    entry->next = next;
}

void* map_compute_if_absent(map_t* map, void* key, void* default_val) {
    struct __map_entry* entry = fetch_entry(map, key);
    if (entry->next != NULL) return entry->value;

    insert_allocating(map, entry, key, default_val);
    return default_val;
}

void map_put(map_t* map, void* key, void* val) {
    struct __map_entry* entry = fetch_entry(map, key);
    if (entry->next != NULL) {
        entry->value = val;
        return;
    }

    insert_allocating(map, entry, key, val);
}

void* map_remove(map_t* map, const void* key) {
    struct __map_entry* entry = fetch_entry(map, key);
    if (entry->next == NULL) return NULL;

    map->size--;
    void* rv = entry->value;
    struct __map_entry* steal = entry->next;
    *entry = *steal;
    free(steal);
    return rv;
}

void map_foreach(map_t* map, foreach_func_t foreach_func, void* data) {
    for (size_t i = 0; i < map->__num_buckets; i++) {
        struct __map_entry* entry = &map->__buckets[i];
        while (entry->next != NULL) {
            foreach_func(entry->key, entry->value, data);
            entry = entry->next;
        }
    }
}