#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <assert.h>
#include "hamt.h"

#include "mempool.h"
MEMPOOL_GENERATE(hamt, struct hamt, 16)
MEMPOOL_GENERATE(hl,   struct hamt_nodelist, 64)
MEMPOOL_GENERATE(hi,   struct hamt_item, 32)

#define tag_ptr(ptr, tag)   ((uintptr_t)(ptr) | (tag))
#define untag_ptr(ptr, tag) ((uintptr_t)(ptr) & ~(tag))
#define is_tagged(ptr, tag) ((uintptr_t)(ptr) & (tag))

#define HAMT_ITEM_TAG 0x1
#define TAG_NODELIST(ptr) \
    (hamtptr_t)untag_ptr(ptr, HAMT_ITEM_TAG)
#define TAG_ITEM(ptr)     \
    (hamtptr_t)tag_ptr(ptr, HAMT_ITEM_TAG)
#define AS_NODELIST(ptr) \
    ((struct hamt_nodelist *)untag_ptr(ptr, HAMT_ITEM_TAG))
#define AS_ITEM(ptr)     \
    ((struct hamt_item *)untag_ptr(ptr, HAMT_ITEM_TAG))
#define AS_VOIDPTR(ptr)  \
    ((void *)untag_ptr(ptr, HAMT_ITEM_TAG)
#define IS_NODELIST(hamtptr) \
    !is_tagged(hamtptr, HAMT_ITEM_TAG)
#define IS_ITEM(hamtptr)     \
    is_tagged(hamtptr, HAMT_ITEM_TAG)

#define for_each_item(item, head)                        \
    for(struct hamt_item *item = (head), *next = NULL;   \
        item && (next = item->next, 1); item = next)

#define popcount(i) __builtin_popcount(i)

#define BITS 6
#define BITS_MASK 0x3f

static inline struct hamt_nodelist *hamt_nodelist_alloc(void);
static inline struct hamt_item *hamt_item_alloc(void);

static void hamtptr_destroy(hamtptr_t hamtptr);
static inline void hamtptr_add_ref(hamtptr_t hamtptr);

// static void hamt_print_hamtptr(hamtptr_t hamtptr, int depth);

hamt_t hamt_create(hamt_equal_fn equal_fn, hamt_hash_fn hash_fn)
{
    // hamt_t hamt = malloc(sizeof(*hamt));
    hamt_t hamt = hamt_mempool_allocate();
    hamt->equal_fn = equal_fn;
    hamt->hash_fn  = hash_fn;
    hamt->root = TAG_NODELIST(hamt_nodelist_alloc());

    return hamt;
}

void hamt_destroy(hamt_t hamt)
{
    if(!hamt) return;

    hamtptr_destroy(hamt->root);
    // free(hamt);
    hamt_mempool_free(hamt);
}

hamt_t hamt_clone(hamt_t src)
{
    // hamt_t hamt = malloc(sizeof(*hamt));
    hamt_t hamt = hamt_mempool_allocate();

    if(!src) return hamt;

    hamt->equal_fn = src->equal_fn;
    hamt->hash_fn  = src->hash_fn;
    hamt->root     = src->root;
    hamtptr_add_ref(hamt->root);

    return hamt;
}

static int hamt_find_hamtptr(hamtptr_t root, hamtptr_t *ret, uint32_t *hash)
{
    *ret = root;

    for(size_t i = 0; i < sizeof(*hash)*8/BITS; i++) {
        assert(IS_NODELIST(*ret));
        struct hamt_nodelist *nodelist = AS_NODELIST(*ret);

        size_t rawidx = *hash & BITS_MASK;
        if(!(nodelist->bitmask & (1 << rawidx))) {
            return sizeof(*hash)*8/BITS - i;
        }

        *hash >>= BITS;

        size_t idx = popcount(nodelist->bitmask & ((1 << rawidx) - 1));
        *ret = nodelist->list[idx];
    }

    return 0;
}

int hamt_get(hamt_t hamt, void *key, void **data)
{
    hamtptr_t hamtptr;
    uint32_t hash = hamt->hash_fn(key);

    if(hamt_find_hamtptr(hamt->root, &hamtptr, &hash) != 0)
        return 1;

    assert(IS_ITEM(hamtptr));
    struct hamt_item *item = AS_ITEM(hamtptr);

    while(!hamt->equal_fn(key, item->key)) {
        if(item->next == NULL) return 1;
        item = item->next;
    }

    *data = item->data;
    return 0;
}

static hamtptr_t hamt_build(uint32_t hash, size_t iter, struct hamt_item **item)
{
      if(iter == 1) {
        *item = hamt_item_alloc();
        return TAG_ITEM(*item);
    }

    hash >>= BITS;
    size_t next_idx = hash & BITS_MASK;

    hamtptr_t next = hamt_build(hash, iter-1, item);

    struct hamt_nodelist *nodelist = hamt_nodelist_alloc();

    // nodelist->list = &next;
    nodelist->list = calloc(1, sizeof(*nodelist->list));
    nodelist->list[0] = next;

    nodelist->bitmask = 1 << next_idx;

    return TAG_NODELIST(nodelist);
}

static hamtptr_t hamt_insert_hamtptr(hamtptr_t root, hamtptr_t hamtptr, uint32_t hash, hamt_equal_fn equal_fn)
{
    if(IS_ITEM(root)) {
        struct hamt_item *item = AS_ITEM(root);

        if(equal_fn(item->key, AS_ITEM(hamtptr)->key)) {
            item->refs--;
            AS_ITEM(hamtptr)->next = item->next;
            return hamtptr;
        }

        if(item->next) {
            hamtptr = hamt_insert_hamtptr(TAG_ITEM(item->next), hamtptr, hash, equal_fn);
        }

        struct hamt_item *new;

        if(item->refs == 1) {
            new = item;
        } else {
            item->refs--;
            new = hamt_item_alloc();
            new->key = item->key;
            new->data = item->data;
        }

        new->next = AS_ITEM(hamtptr);
        return TAG_ITEM(new);
    }

    struct hamt_nodelist *nodelist = AS_NODELIST(root);

    size_t rawidx = hash & BITS_MASK;
    size_t idx = popcount(nodelist->bitmask & ((1 << rawidx) - 1));

    size_t list_len = popcount(nodelist->bitmask);
    size_t newlist_len = list_len;

    if(nodelist->bitmask & (1 << rawidx)) {
        hamtptr = hamt_insert_hamtptr(nodelist->list[idx], hamtptr, hash >> BITS, equal_fn);

        if(nodelist->refs == 1) {
            nodelist->list[idx] = hamtptr;
            return TAG_NODELIST(nodelist);
        }
    } else {
        newlist_len++;
    }

    hamtptr_t *newlist = calloc(newlist_len, sizeof(*newlist));
    newlist[idx] = hamtptr;

    for(size_t i = 0; i < idx; i++)
        newlist[i] = nodelist->list[i];

    if(list_len == newlist_len)
        for(size_t i = idx+1; i < newlist_len; i++)
            newlist[i] = nodelist->list[i];
    else
        for(size_t i = idx+1; i < newlist_len; i++)
            newlist[i] = nodelist->list[i-1];

    struct hamt_nodelist *new;

    if(nodelist->refs == 1) {
        new = nodelist;
        free(nodelist->list);
    } else {
        nodelist->refs--;
        new = hamt_nodelist_alloc();
    }

    new->list = newlist;
    new->bitmask = nodelist->bitmask | (1 << rawidx);
    return TAG_NODELIST(new);
}

int hamt_set(hamt_t hamt, void *key, void *data, void **keyptr, void **prevdata)
{
    hamtptr_t hamtptr;
    uint32_t hash = hamt->hash_fn(key);
    uint32_t hash_cpy = hash;

    size_t iter;
    if((iter = hamt_find_hamtptr(hamt->root, &hamtptr, &hash_cpy)) != 0) {
        struct hamt_item *item;
        hamtptr_t new = hamt_build(hash_cpy, iter, &item);

        hamt->root = hamt_insert_hamtptr(hamt->root, new, hash, hamt->equal_fn);

        item->key = key;
        item->data = data;
        return 0;
    }

    assert(IS_ITEM(hamtptr));

    for_each_item(item, AS_ITEM(hamtptr)) {
        if(hamt->equal_fn(item->key, key)) {
            if(keyptr)   *keyptr   = item->key;
            if(prevdata) *prevdata = item->data;
            if(item->refs == 1) {
                item->data = data;
                return 0;
            }

            key = item->key;
            break;
        }
    }

    struct hamt_item *new = hamt_item_alloc();
    new->key = key;
    new->data = data;

    hamt->root = hamt_insert_hamtptr(hamt->root, TAG_ITEM(new), hash, hamt->equal_fn);
    return 0;
}

static inline struct hamt_nodelist *hamt_nodelist_alloc(void)
{
    // struct hamt_nodelist *nodelist = malloc(sizeof(*nodelist));
    struct hamt_nodelist *nodelist = hl_mempool_allocate();
    nodelist->refs = 1;
    nodelist->list = 0;
    nodelist->bitmask = 0;
    return nodelist;
}

static inline struct hamt_item *hamt_item_alloc(void)
{
    // struct hamt_item *item = malloc(sizeof(*item));
    struct hamt_item *item = hi_mempool_allocate();
    item->refs = 1;
    item->key = NULL;
    item->data = NULL;
    item->next = NULL;
    return item;
}

static inline void hamtptr_destroy(hamtptr_t hamtptr)
{
    if(IS_NODELIST(hamtptr)) {
        struct hamt_nodelist *nodelist = AS_NODELIST(hamtptr);

        for(size_t i = 0; i < popcount(nodelist->bitmask); i++)
            hamtptr_destroy(nodelist->list[i]);

        if(--nodelist->refs == 0) {
            free(nodelist->list);
            // free(nodelist);
            hl_mempool_free(nodelist);
        }
    } else {
        for_each_item(item, AS_ITEM(hamtptr))
            if(--item->refs == 0)
                hi_mempool_free(item);
                // free(item);
    }
}

static inline void hamtptr_add_ref(hamtptr_t hamtptr)
{
    if(IS_NODELIST(hamtptr)) {
        for(size_t i = 0; i < popcount(AS_NODELIST(hamtptr)->bitmask); i++)
            hamtptr_add_ref(AS_NODELIST(hamtptr)->list[i]);
        AS_NODELIST(hamtptr)->refs++;
    } else {
        for_each_item(item, AS_ITEM(hamtptr))
            item->refs++;
    }
}

// static void hamt_print_hamtptr(hamtptr_t hamtptr, int depth)
// {
//     for(int i = 0; i < depth; i++) printf("  ");

//     if(IS_NODELIST(hamtptr)) {
//         printf("%d, MASK %b\n", AS_NODELIST(hamtptr)->refs, AS_NODELIST(hamtptr)->bitmask);
//         for(size_t i = 0; i < popcount(AS_NODELIST(hamtptr)->bitmask); i++) {
//             hamt_print_hamtptr(AS_NODELIST(hamtptr)->list[i], depth+1);
//         }
//     } else {
//         printf("%d, %s: %s\n", AS_ITEM(hamtptr)->refs, (char *)AS_ITEM(hamtptr)->key, (char *)AS_ITEM(hamtptr)->data);
//         if(AS_ITEM(hamtptr)->next)
//             hamt_print_hamtptr(TAG_ITEM(AS_ITEM(hamtptr)->next), depth+1);
//     }
// }