#include "collections/HashMap.h"
#include <assert.h>

//TODO: sort bucket keys for binary search

#define __HashMap_HASH_FUNC str_hash32
#define __HashMapBucket_MAX_LEN 16

static const Array __HashMap_heights = Array_CONST(u32, {
    17, 31, 61, 127, 257, 521, 1021, 2053, 4099, 8191, 16381, 32771,
    65521, 131071, 262147, 524287, 1048583, 2097169, 4194319,
    8388617, 16777213, 33554467, 67108859, 134217757, 268435493
});

void HashMap_alloc(HashMap* ptr, u32 value_t_size, FreeFunction NULLABLE(value_destructor)){
    ptr->value_t_size = value_t_size;
    ptr->value_destructor = value_destructor;
    ptr->height_n = 0;
    ptr->height = ((u32*)__HashMap_heights.data)[0];
    u32 alloc_size = ptr->height * sizeof(HashMapBucket);
    ptr->table = (HashMapBucket*)malloc(alloc_size);
    memset(ptr->table, 0, alloc_size);
}

void HashMap_free(HashMap* ptr){
    for(u32 i = 0; i < ptr->height; i++){
        HashMapBucket* bu = &ptr->table[i];
        u32 len = List_len(&bu->key_hash_list, KeyHash);
        
        // free key strings
        for(u32 j = 0; j < len; j++){
            KeyHash* kh = (KeyHash*)bu->key_hash_list.data + j;
            free(kh->key.data);
        }
        
        // destroy values
        if(ptr->value_destructor){
            u8* value_ptr = (u8*)bu->value_list.data;
            u8* end = value_ptr + bu->value_list.size;
            while(value_ptr < end){
                ptr->value_destructor(value_ptr);
                value_ptr += ptr->value_t_size;
            }
        }

        free(bu->key_hash_list.data);
        free(bu->value_list.data);
    }
    
    free(ptr->table);
}

typedef struct BucketAndIndex {
    HashMapBucket* bu;
    i32 i; // -1 if not found
} BucketAndIndex;

#define BucketAndIndex_null ((BucketAndIndex){ .bu = NULL, .i = -1 })

static BucketAndIndex __HashMap_search(HashMap* ptr, str key, u32 hash){
    BucketAndIndex r;
    r.bu = &ptr->table[hash % ptr->height];
    for(r.i = 0; r.i < (i32)List_len(&r.bu->key_hash_list, KeyHash); r.i++){
        KeyHash* kh = (KeyHash*)r.bu->key_hash_list.data + r.i;
        if(kh->hash == hash && str_equals(kh->key, key)){
            return r;
        }
    }

    r.i = -1;
    return r;
}

void* NULLABLE(HashMap_tryGetPtr)(HashMap* ptr, str key){
    u32 hash = __HashMap_HASH_FUNC(key);
    BucketAndIndex r = __HashMap_search(ptr, key, hash);
    // key not found
    if(r.i == -1)
        return NULL;
    return ((u8*)r.bu->value_list.data) + r.i * ptr->value_t_size;
}


static void __HashMap_expand(HashMap* ptr){
    u32 height_expanded_n = ptr->height_n + 1;
    assert(height_expanded_n >= Array_len(&__HashMap_heights, u32) && "HashMap IS FULL! Fix your code.");

    // alloc new HashMapBucket array
    u32 height_expanded = ((u32*)__HashMap_heights.data)[height_expanded_n];
    u32 table_expanded_size = height_expanded * sizeof(HashMapBucket);
    HashMapBucket* table_expanded = (HashMapBucket*)malloc(table_expanded_size);
    memset(table_expanded, 0, table_expanded_size);

    // copy values from old buckets to new
    for(u32 i = 0; i < ptr->height; i++){
        HashMapBucket* old_bucket = &ptr->table[i];
        u32 len = List_len(&old_bucket->key_hash_list, KeyHash);
        
        for(u32 j = 0; j < len; j++){
            KeyHash kh = ((KeyHash*)old_bucket->key_hash_list.data)[j];
            HashMapBucket* new_bucket = &table_expanded[kh.hash % height_expanded];
            List_push(&new_bucket->key_hash_list, KeyHash, kh);
            void* old_value_ptr = (u8*)old_bucket->value_list.data + j * ptr->value_t_size;
            List_push_size(&new_bucket->value_list, old_value_ptr, ptr->value_t_size);
        }
        
        free(old_bucket->key_hash_list.data);
        free(old_bucket->value_list.data);
    }
    free(ptr->table);
    ptr->table = table_expanded;
    ptr->height = height_expanded;
    ptr->height_n = height_expanded_n;
}

bool HashMap_tryPush(HashMap* ptr, str key, void* value_ptr){
    u32 hash = __HashMap_HASH_FUNC(key);
    BucketAndIndex r = __HashMap_search(ptr, key, hash);
    // found existing item with the same key
    if(r.i != -1)
        return false;

    HashMapBucket* bu = r.bu;
    if(List_len(&bu->key_hash_list, KeyHash) >= __HashMapBucket_MAX_LEN){
        __HashMap_expand(ptr);
        bu = &ptr->table[hash % ptr->height];
    }
    
    KeyHash kh = { .key = str_copy(key), .hash = hash };
    List_push(&bu->key_hash_list, KeyHash, kh);
    List_push_size(&bu->value_list, value_ptr, ptr->value_t_size);
    return true;
}

bool HashMap_tryDelete(HashMap* ptr, str key){
    u32 hash = __HashMap_HASH_FUNC(key);
    BucketAndIndex r = __HashMap_search(ptr, key, hash);
    // key not found
    if(r.i == -1)
        return false;
    
    List_removeAt(&r.bu->key_hash_list, KeyHash, r.i, 1);
    List_removeAt_size(&r.bu->value_list, r.i, ptr->value_t_size);
    return true;
}