#include "../../include/std/hash_map.h"
#include <assert.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

typedef struct TslHashMapNode TslHashMapNode;
struct TslHashMapNode {
    void *data;
    TslHashMapNode *next;
};

static void hash_map_node_get(TslHashMapNode *self, uint64_t *hash, TslStringView *key, size_t *size, uint8_t **data) {
    memcpy(hash, (uint8_t*)self->data, sizeof(uint64_t));
    memcpy(&key->size, (uint8_t*)self->data + sizeof(uint64_t), sizeof(key->size));
    key->data = (const char*)self->data + sizeof(uint64_t) + sizeof(key->size);
    memcpy(size, (uint8_t*)self->data + sizeof(uint64_t) + sizeof(key->size) + key->size, sizeof(size_t));
    *data = (uint8_t*)self->data + sizeof(uint64_t) + sizeof(key->size) + key->size + sizeof(size_t);
}

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

void tsl_hash_map_deinit(TslHashMap *self) {
    TslHashMapNode **bucket = (TslHashMapNode**)self->buckets_data;
    TslHashMapNode **bucket_end = (TslHashMapNode**)((char*)self->buckets_data + self->buckets_capacity);
    while(bucket != bucket_end) {
        TslHashMapNode *node = *bucket;
        while(node) {
            TslHashMapNode *prev_node = node;
            node = node->next;
            free(prev_node->data);
            free(prev_node);
        }
        ++bucket;
    }
    free(self->buckets_data);
}

/* TODO: Remove if (data) and if (output) */
static int tsl_hash_map_append_bucket(TslHashMapNode **head_node, uint64_t hash, const TslStringView *key, size_t size, const void *data, void **output) {
    TslHashMapNode *next_node;
    uint8_t *node_data = malloc(sizeof(hash) + sizeof(key->size) + key->size + sizeof(size) + size);
    if(output)
        *output = node_data;
    if(!node_data) {
        fprintf(stderr, "Error: hash map append failed. Reason: out of memory\n");
        return 0;
    }

    next_node = malloc(sizeof(TslHashMapNode));
    if(!next_node) {
        free(node_data);
        if(output)
            *output = NULL;
        fprintf(stderr, "Error: hash map append failed. Reason: out of memory\n");
        return 0;
    }

    memcpy(node_data, &hash, sizeof(hash));
    /* TODO: Instead of allocating space for the key, use the key data pointer and size directly. */
    memcpy(node_data + sizeof(hash), &key->size, sizeof(key->size));
    memcpy(node_data + sizeof(hash) + sizeof(key->size), key->data, key->size);
    /*
        TODO: Instead of allocating space for the data, allow the user to pass a pointer in the insert
        method and use that directly.
    */
    memcpy(node_data + sizeof(hash) + sizeof(key->size) + key->size, &size, sizeof(size));
    if(data)
        memcpy(node_data + sizeof(hash) + sizeof(key->size) + key->size + sizeof(size), data, size);

    next_node->data = node_data;
    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) {
            uint64_t hash;
            TslStringView key;
            size_t size;
            uint8_t *data;
            size_t index;
            hash_map_node_get(node, &hash, &key, &size, &data);

            index = tsl_hash_map_get_index(self, 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 = 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;
    {
        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;
        }
    }

    if(old_capacity == 0)
        return 1;

    tsl_hash_map_reorder_nodes(self, old_capacity);
    return 1;
}

int tsl_hash_map_insert(TslHashMap *self, const TslStringView *key, const void *data, size_t size, TslHashFunc hash_func) {
    uint64_t hash = hash_func(key->data, key->size);
    size_t index;
    TslHashMapNode **bucket;
    assert(!tsl_hash_map_get(self, key, hash_func));
    assert(size > 0);
    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, size, data, NULL);
}

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

    hash = hash_func(key->data, key->size);
    index = tsl_hash_map_get_index(self, hash);
    node = *((TslHashMapNode**)self->buckets_data + index);
    while(node) {
        uint64_t node_hash;
        TslStringView node_key;
        size_t node_size;
        uint8_t *node_data;
        hash_map_node_get(node, &node_hash, &node_key, &node_size, &node_data);
        if(hash == node_hash && key->size == node_key.size && memcmp(key->data, node_key.data, node_key.size) == 0)
            return node_data;
        node = node->next;
    }

    return NULL;
}

int tsl_hash_map_get_or_create(TslHashMap *self, const TslStringView *key, size_t size, TslHashFunc hash_func, void **output) {
    uint64_t hash;
    size_t index;
    TslHashMapNode **bucket;
    TslHashMapNode *node;
    assert(size > 0);
    if(self->buckets_capacity == 0) {
        if(!tsl_hash_map_ensure_bucket_capacity_for_one_new_item(self)) {
            *output = NULL;
            return 0;
        }
    }

    hash = hash_func(key->data, key->size);
    index = tsl_hash_map_get_index(self, hash);
    bucket = (TslHashMapNode**)self->buckets_data + index;
    node = *bucket;
    while(node) {
        uint64_t node_hash;
        TslStringView node_key;
        size_t node_size;
        uint8_t *node_data;
        hash_map_node_get(node, &node_hash, &node_key, &node_size, &node_data);
        if(hash == node_hash && key->size == node_key.size && memcmp(key->data, node_key.data, node_key.size) == 0) {
            *output = node_data;
            return 1;
        }
        node = node->next;
    }

    return tsl_hash_map_append_bucket(bucket, hash, key, size, NULL, output);
}