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

#define DEFAULT_CAPACITY 16
#define TOMBSTONE_VAL ((void*)-1)

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;
}

/* MAPS */

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", map->__num_buckets);
}

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_map_slot(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_map_slot(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_map_slot(map, key);
    return entry == NULL ? default_val : entry->value;
}

static void rehash_map(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_map_slot(map, entry->key);
            *slot = entry;
            entry = next;
        }
    }

    free(old_buckets);
}

static void map_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(map);
        slot = fetch_map_slot(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_map_slot(map, key);
    struct __map_entry* entry = *slot;
    if (entry != NULL) return entry->value;

    map_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_map_slot(map, key);
    struct __map_entry* entry = *slot;
    if (entry != NULL) {
        entry->value = val;
        return;
    }

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

void* map_remove(map_t* map, const void* key) {
    struct __map_entry** slot = fetch_map_slot(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,
    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 != NULL) {
            foreach_func(entry->key, entry->value, data);
            entry = entry->next;
        }
    }
}

void map_foreach_readwrite(
    map_t* map,
    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);
}

/* SETS */

void set_init(
    set_t* set,
    hash_func_t hash_func,
    eq_func_t eq_func,
    double load_limit
) {
    set_init_capacity(set, hash_func, eq_func, load_limit, DEFAULT_CAPACITY);
}

void set_init_capacity(
    set_t* set,
    hash_func_t hash_func,
    eq_func_t eq_func,
    double load_limit,
    size_t capacity
) {
    set->hash_func = hash_func;
    set->eq_func = eq_func;
    set->load_limit = load_limit;
    set->size = 0;
    set->__num_buckets = capacity / load_limit + 1;
    set->__buckets = calloc(set->__num_buckets, sizeof(void*));

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

void set_destroy(set_t* set) {
    free(set->__buckets);
}

static size_t fetch_set_idx(
    const set_t* set,
    const void* key,
    bool fetch_tombstones
) {
    if (key == NULL || key == TOMBSTONE_VAL)
        crash("safec sets do not support non-address keys: %p.\n", key);
    size_t base = set->hash_func(key) % set->__num_buckets;
    size_t idx = base, i = 0, offset = 0;

    for (;;) {
        void* cur = set->__buckets[idx];
        if (cur == NULL || (fetch_tombstones && cur == TOMBSTONE_VAL)
                || (cur != TOMBSTONE_VAL && set->eq_func(key, cur)))
            return idx;
        if (offset >= set->__num_buckets) return set->__num_buckets;
        i++;
        offset = i * i;
        idx = (base + offset) % set->__num_buckets;
    }
}

bool set_contains(const set_t* set, const void* key) {
    size_t idx = fetch_set_idx(set, key, false);
    return idx < set->__num_buckets && set->__buckets[idx] != NULL;
}

void* set_get(const set_t* set, const void* key) {
    return set_get_or_default(set, key, NULL);
}

void* set_get_or_default(const set_t* set, const void* key, void* default_key) {
    size_t idx = fetch_set_idx(set, key, false);
    if (idx >= set->__num_buckets) return default_key;
    void* found = set->__buckets[idx];
    return found == NULL ? default_key : found;
}

static void rehash_set(set_t* set) {
    size_t old_num_buckets = set->__num_buckets;
    void** old_buckets = set->__buckets;
    set->__num_buckets <<= 1;
    set->__buckets = calloc(set->__num_buckets, sizeof(void*));
    if (set->__buckets == NULL)
        crash(
            "Out of memory allocating %ld hash buckets\n", set->__num_buckets);

    for (size_t i = 0; i < old_num_buckets; i++) {
        void* key = old_buckets[i];
        if (key == NULL) continue;
        size_t idx = fetch_set_idx(set, key, true);
        if (idx >= set->__num_buckets)
            crash(
                "Set failed rehashing, likely due to bad hash function.\n");
        set->__buckets[idx] = key;
    }
    free(old_buckets);
}

static size_t fetch_set_idx_rehashing(set_t* set, void* key) {
    size_t idx = fetch_set_idx(set, key, true);
    if ((double) ++set->size / set->__num_buckets > set->load_limit
            || idx >= set->__num_buckets) {
        rehash_set(set);
        idx = fetch_set_idx(set, key, true);
        if (idx >= set->__num_buckets)
            crash(
                "Set still full after rehashing, "
                "likely due to bad hash function.\n");
    }
    return idx;
}

void* set_add(set_t* set, void* key) {
    size_t idx = fetch_set_idx_rehashing(set, key);
    void* old_key = set->__buckets[idx];
    if (old_key != NULL) return old_key;
    set->__buckets[idx] = key;
    return key;
}

void set_put(set_t* set, void* key) {
    size_t idx = fetch_set_idx_rehashing(set, key);
    set->__buckets[idx] = key;
}

void* set_remove(set_t* set, const void* key) {
    size_t idx = fetch_set_idx(set, key, false);
    void* old_key = set->__buckets[idx];
    if (old_key != NULL) set->size--;
    set->__buckets[idx] = TOMBSTONE_VAL;
    return old_key;   
}

void set_foreach(set_t* set, set_foreach_func_t foreach_func, void* data) {
    for (size_t i = 0; i < set->__num_buckets; i++) {
        void* key = set->__buckets[i];
        if (key == NULL || key == TOMBSTONE_VAL) continue;
        foreach_func(key, data);
    }
}