#include "../../include/std_gc/hash_map.h"
#include <assert.h>
#include <string.h>
#include <stdio.h>
#include <gc.h>

void tsl_hash_map_init(TslHashMap *self) {
    self->buckets_data = NULL;
    self->buckets_size = 0;
    self->buckets_capacity = 0;
}

/* TODO: Remove if (data) and if (output) */
static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash, const TslValue *key, TslValue *value, TslValue **output) {
    TslHashMapNode *next_node;

    if(output)
        *output = NULL;

    next_node = GC_MALLOC(sizeof(TslHashMapNode));
    if(!next_node) {
        fprintf(stderr, "Error: hash map insert failed. Reason: out of memory\n");
        return 0;
    }

    next_node->hash = hash;
    next_node->key = *key;
    if(value)
        next_node->value = *value;

    if(output)
        *output = &next_node->value;

    if(*head_node) {
        next_node->next = (*head_node)->next;
        (*head_node)->next = next_node;
    } else {
        next_node->next = NULL;
        *head_node = next_node;
    }

    return 1;
}

static size_t tsl_hash_map_get_index(TslHashMap *self, uint64_t hash) {
    assert((self->buckets_capacity >> 3) > 0);
    /* >>3 = 8 = sizeof(TslHashMapNode*) */
    return hash & ((self->buckets_capacity >> 3) - 1);
}

static void tsl_hash_map_reorder_nodes(TslHashMap *self, size_t old_capacity) {
    TslHashMapNode **bucket = (TslHashMapNode**)self->buckets_data;
    TslHashMapNode **bucket_end = (TslHashMapNode**)((char*)self->buckets_data + old_capacity);
    size_t bucket_index = 0;
    while(bucket != bucket_end) {
        TslHashMapNode *node = *bucket;
        TslHashMapNode *prev_node = node; /* Set to node for optimization reason, where prev_node->next = node->next; which becomes no-op */
        int all_nodes_moved = 1;

        while(node) {
            size_t index = tsl_hash_map_get_index(self, node->hash);
            if(index != bucket_index) {
                TslHashMapNode **new_bucket = (TslHashMapNode**)self->buckets_data + index;
                prev_node->next = node->next;
                if(*new_bucket) {
                    node->next = (*new_bucket)->next;
                    (*new_bucket)->next = node;
                } else {
                    node->next = NULL;
                    *new_bucket = node;
                }
            } else {
                all_nodes_moved = 0;
            }

            prev_node = node;
            node = node->next;
        }

        if(all_nodes_moved)
            *bucket = NULL;

        ++bucket_index;
        ++bucket;
    }
}

static int tsl_hash_map_ensure_bucket_capacity_for_one_new_item(TslHashMap *self) {
    void *new_ptr;
    size_t new_capacity = self->buckets_capacity;
    size_t new_size = self->buckets_size + sizeof(TslHashMapNode*);
    size_t old_capacity = self->buckets_capacity;
    self->buckets_size = new_size;
    if(new_size <= self->buckets_capacity)
        return 1;

    if(new_capacity == 0)
        new_capacity = sizeof(TslHashMapNode*) << 3;

    while(new_capacity < new_size) {
        new_capacity <<= 1;
    }

    new_ptr = GC_REALLOC(self->buckets_data, new_capacity);
    if(!new_ptr) {
        fprintf(stderr, "Error: hash map realloc failed. Reason: out of memory\n");
        return 0;
    }

    self->buckets_data = new_ptr;
    self->buckets_capacity = new_capacity;
    {
        /* TODO: Remove this. This is not needed since GC_REALLOC sets the data to 0 */
        TslHashMapNode **bucket = (TslHashMapNode**)((char*)self->buckets_data + old_capacity);
        TslHashMapNode **bucket_end = (TslHashMapNode**)((char*)self->buckets_data + new_capacity);
        while(bucket != bucket_end) {
            *bucket = NULL;
            ++bucket;
        }
    }

    tsl_hash_map_reorder_nodes(self, old_capacity);
    return 1;
}

int tsl_hash_map_insert(TslHashMap *self, const TslValue *key, TslValue *value) {
    uint64_t hash = tsl_value_hash(key);
    size_t index;
    TslHashMapNode **bucket;
    assert(!tsl_hash_map_get(self, key));
    if(!tsl_hash_map_ensure_bucket_capacity_for_one_new_item(self))
        return 0;

    index = tsl_hash_map_get_index(self, hash);
    bucket = (TslHashMapNode**)self->buckets_data + index;
    return tsl_hash_map_append_bucket(bucket, hash, key, value, NULL);
}

TslValue* tsl_hash_map_get(TslHashMap *self, const TslValue *key) {
    uint64_t hash;
    size_t index;
    TslHashMapNode *node;
    if(self->buckets_capacity == 0)
        return NULL;

    hash = tsl_value_hash(key);
    index = tsl_hash_map_get_index(self, hash);
    node = *((TslHashMapNode**)self->buckets_data + index);
    while(node) {
        if(hash == node->hash && tsl_value_equals(key, &node->key))
            return &node->value;
        node = node->next;
    }

    return NULL;
}

TslValue* tsl_hash_map_get_or_create(TslHashMap *self, const TslValue *key) {
    uint64_t hash;
    size_t index;
    TslValue *value;
    TslHashMapNode **bucket;

    hash = tsl_value_hash(key);
    if(self->buckets_capacity > 0) {
        size_t index = tsl_hash_map_get_index(self, hash);
        TslHashMapNode **bucket = (TslHashMapNode**)self->buckets_data + index;
        TslHashMapNode *node = *bucket;
        while(node) {
            if(hash == node->hash && tsl_value_equals(key, &node->key))
                return &node->value;
            node = node->next;
        }
    }

    if(!tsl_hash_map_ensure_bucket_capacity_for_one_new_item(self))
        return NULL;

    index = tsl_hash_map_get_index(self, hash);
    bucket = (TslHashMapNode**)self->buckets_data + index;

    if(!tsl_hash_map_append_bucket(bucket, hash, key, NULL, &value))
        return NULL;
    
    return value;
}

void tsl_hash_map_create_iterator(TslHashMap *self, TslHashMapIterator *iterator) {
    iterator->hash_map = self;
    iterator->index = -1;
    iterator->node = NULL;
}

int tsl_hash_map_iterator_next(TslHashMapIterator *self) {
    if(self->node) {
        TslHashMapNode *new_node = self->node->next;
        if(new_node) {
            self->node = new_node;
            return 1;
        }
    }

    ++self->index;
    while(self->index < self->hash_map->buckets_size) {
        TslHashMapNode **bucket = (TslHashMapNode**)self->hash_map->buckets_data + self->index;
        if(*bucket) {
            self->node = *bucket;
            return 1;
        }
        ++self->index;
    }

    return 0;
}