#include <c_HashST.h>

#define DEFAULT_INITIAL_CAPACITY (1024*8)
#define GROWTH_FACTOR            1

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

static c_HashSTNode_t* GetNode(c_HashST_t* self, const void* key) {
    c_HashSTNode_t* result = NULL;
    c_HashSTNode_t* node = NULL;
    c_PtrBag_t* bucket = NULL;

    uint32_t hash=0;
    c_size_t bucket_idx = 0;

    hash = self->key_ops.hash(key, self->key_ops.arg);
    bucket_idx = hash % self->capacity;
    bucket = self->buckets[bucket_idx];

    if (!bucket) {
        return NULL;
    }

    for (c_size_t i=0; i<bucket->size; i++) {
        node = c_PtrBag_Get(bucket, i);
        if (!node) continue;
        if (node->hash == hash) {
            if (self->key_ops.eq(node->key, key, self->key_ops.arg)) {
                result = node;
                break;
            }
        }
    }

    return result;
}


C_STATIC_FORCE_INLINE
c_size_t NumCol(c_PtrBag_t* self) {
    if (!self) return 0;
    return (self->size==0)?0:self->size-1;
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


c_err_t c_HashST_Init(c_HashST_t* self, c_size_t capacity, c_HashKeyOps_t key_ops, c_HashValOps_t val_ops, c_Allocator_t* allocator) {
    if (!self ) return C_ERR_PARAM;

    self->capacity = capacity==0?DEFAULT_INITIAL_CAPACITY:capacity;
    self->key_ops = key_ops;
    self->val_ops = val_ops;
    self->allocator = allocator!=NULL?*allocator:c_DefaultAllocator;
    self->size = 0;

    // key_ops.arg = &self->allocator;
    // val_ops.arg = &self->allocator;

    self->buckets = c_Allocator_Alloc(&self->allocator, self->capacity * sizeof(*(self->buckets)));
    if (!self->buckets) return C_ERR_NOMEM;

    for (c_size_t i=0; i<self->capacity; i++) {
        self->buckets[i] = NULL;
    }

    return C_ERR_OK;
}

void c_HashST_Destroy(c_HashST_t* self) {
    if (!self ) return;
    c_PtrBag_t* bucket=0;
    c_HashSTNode_t* node=0;

    for (c_size_t i=0; i<self->capacity; i++) {
        bucket = self->buckets[i];
        if (!bucket) continue;
        for (c_size_t j=0; j<bucket->size; j++) {
            node = c_PtrBag_Get(bucket, j);
            if (!node) continue;
            self->key_ops.free(node->key, self->key_ops.arg);
            self->val_ops.free(node->val, self->val_ops.arg);
            c_Allocator_Free(&self->allocator, node);
        }
        c_PtrBag_Destroy(bucket);
        c_Allocator_Free(&self->allocator, bucket);
    }

    c_Allocator_Free(&self->allocator, self->buckets);
}

void* c_HashST_Get(c_HashST_t* self, const void* key) {
    c_HashSTNode_t* node = GetNode(self, key);
    if (!node) return NULL;
    return node->val;
}

c_err_t c_HashST_Resize(c_HashST_t* self, c_size_t new_capacity) {
    if (!self || new_capacity < self->capacity) return C_ERR_PARAM;
    if (new_capacity == self->capacity) {
        return C_ERR_OK;
    }

    c_PtrBag_t** new_buckets = c_Allocator_Alloc(&self->allocator, new_capacity * sizeof(*(self->buckets)));
    if (!new_buckets) return C_ERR_NOMEM;
    for (c_size_t i=0; i<new_capacity; i++) {
        new_buckets[i] = NULL;
    }

    for (c_size_t i=0; i<self->capacity; i++) {
        c_PtrBag_t* bucket = self->buckets[i];
        if (!bucket) continue;
        for (c_size_t j=0; j<bucket->size; j++) {
            c_HashSTNode_t* node = c_PtrBag_Get(bucket, j);
            if (!node) continue;
            const c_size_t idx = node->hash % new_capacity;
            if (!new_buckets[idx]) {
                new_buckets[idx] = c_Allocator_Alloc(&self->allocator, sizeof(*bucket));
                if (!new_buckets[idx]) return C_ERR_NOMEM;
                c_PtrBag_Init(new_buckets[idx], 0, &self->allocator);
            }
            c_PtrBag_Add(new_buckets[idx], node);
        }
        c_PtrBag_Destroy(bucket);
        c_Allocator_Free(&self->allocator, bucket);
    }

    self->capacity = new_capacity;
    c_Allocator_Free(&self->allocator, self->buckets);
    self->buckets = new_buckets;

    return C_ERR_OK;
}

c_err_t c_HashST_Put(c_HashST_t* self, const void* key, const void* val) {
    if (!self || key==NULL || val==NULL) return C_ERR_PARAM;
    c_HashSTNode_t* node = GetNode(self, key);

    // 情况1: key 已经存在
    if (node!=NULL) {
        self->val_ops.free(node->val, self->val_ops.arg);
        node->val = val?self->val_ops.cp(val, self->val_ops.arg):NULL;
        return C_ERR_OK;
    }

    node = c_Allocator_Alloc(&self->allocator, sizeof(*node));
    if (!node) return C_ERR_NOMEM;
    node->hash = self->key_ops.hash(key, self->key_ops.arg);
    node->key = self->key_ops.cp(key, self->key_ops.arg);
    node->val = self->val_ops.cp(val, self->val_ops.arg);

    const c_size_t idx = node->hash % self->capacity;
    c_PtrBag_t* bucket = self->buckets[idx];
    if (bucket==NULL) {
        bucket = c_Allocator_Alloc(&self->allocator, sizeof(*bucket));
        if (!bucket) return C_ERR_NOMEM;
        c_PtrBag_Init(bucket, 0, &self->allocator);
        self->buckets[idx] = bucket;
    }

    c_PtrBag_Add(bucket, node);

    self->size++;

    // if (self->size > self->capacity * GROWTH_FACTOR) {
    //     return c_HashST_Resize(self, self->capacity * 2);
    // }

    return C_ERR_OK;
}


c_err_t c_HashST_Remove(c_HashST_t* self, const void* key) {
    if (!self || key==NULL) return C_ERR_PARAM;

    c_HashSTNode_t* node = NULL;
    c_PtrBag_t* bucket = NULL;

    uint32_t hash=0;
    c_size_t bucket_idx = 0;

    hash = self->key_ops.hash(key, self->key_ops.arg);
    bucket_idx = hash % self->capacity;
    bucket = self->buckets[bucket_idx];

    if (!bucket) {
        return C_ERR_NOTFOUND;
    }

    for (c_size_t i=0; i<bucket->size; i++) {
        node = c_PtrBag_Get(bucket, i);
        if (!node) continue;
        if (node->hash == hash) {
            if (self->key_ops.eq(node->key, key, self->key_ops.arg)) {
                self->key_ops.free(node->key, self->key_ops.arg);
                self->val_ops.free(node->val, self->val_ops.arg);
                c_Allocator_Free(&self->allocator, node);
                c_PtrBag_RemoveAt(bucket, i, 0);
                self->size--;
                return C_ERR_OK;
            }
        }
    }

    return C_ERR_NOTFOUND;
}

bool c_HashST_Contains(c_HashST_t* self, const void* key) {
    if (!self || key==NULL) return false;
    const c_HashSTNode_t* node = GetNode(self, key);
    if (node==NULL) return false;
    return true;
}


c_size_t c_HashST_NumCol(c_HashST_t* self) {
    if (!self) return 0;
    c_size_t result = 0;
    for (c_size_t i=0; i<self->capacity; i++) {
        result+= NumCol(self->buckets[i]);
    }
    return result;
}

void c_HashST_ForEach(c_HashST_t* self, void (*apply)(c_HashSTNode_t* node, void* args), void* args) {
    if (!self || !apply) return;
    if (self->size==0) return;

    for (c_size_t i=0; i<self->capacity; i++) {
        c_PtrBag_t* bucket = self->buckets[i];
        if (!bucket) continue;
        for (c_size_t j=0; j<bucket->size; j++) {
            c_HashSTNode_t* node = c_PtrBag_Get(bucket, j);
            if (!node) continue;
            apply(node, args);
        }
    }
}