#include "map.h"
#include "crash.h"
#include <stdlib.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];
        while (node != NULL) {
            struct __map_entry* next = node->next;
            free(node);
            node = next;
        }
    }

    free(map->__buckets);
}

static struct __map_entry** fetch_entry(const map_t* map, const void* key) {
    struct __map_entry** bucket =
        &map->__buckets[map->hash_func(key) % map->__num_buckets];
    struct __map_entry* cur = *bucket;
    if (cur == NULL || map->eq_func(key, cur->key)) return bucket;

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

bool map_contains(const map_t* map, const void *key) {
    return *fetch_entry(map, key) != 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 == NULL ? default_val : entry->value;
}

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* entry = old_buckets[i];
        while (entry != NULL) {
            struct __map_entry* next = entry->next;
            entry->next = NULL;
            struct __map_entry** slot = fetch_entry(map, entry->key);
            *slot = entry;
            entry = next;
        }
    }

    free(old_buckets);
}

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

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

    entry->key = key;
    entry->value = val;
    entry->next = NULL;
    *slot = entry;
}

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

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

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

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

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

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

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

void map_foreach_readwrite(
    map_t* map,
    foreach_func_t foreach_func,
    void* data
) {
    void** keys = malloc(map->size * sizeof(void*));
    void** vals = malloc(map->size * sizeof(void*));
    size_t idx = 0;

    for (size_t i = 0; i < map->__num_buckets; i++) {
        struct __map_entry* entry = map->__buckets[i];
        while (entry != NULL) {
            keys[idx] = entry->key;
            vals[idx++] = entry->value;
            entry = entry->next;
        }
    }

    for (size_t i = 0; i < idx; i++) {
        foreach_func(keys[i], vals[i], data);
    }

    free(keys);
    free(vals);
}