#include <assert.h>
#include <stdlib.h>
#include <string.h>

#include "alloc.h"
#include "hashmap.h"

struct apfl_hashmap_bucket {
    void *keys;
    void *values;
    size_t len;

    // These will be the same, unless a memory reallocation failed
    size_t keys_cap;
    size_t values_cap;
};

#define FNV_PRIME 1099511628211U

apfl_hash
apfl_hash_fnv1a_add(const void *data, size_t len, apfl_hash hash)
{
    for (size_t i = 0; i < len; i++) {
        uint8_t byte = ((uint8_t *)data)[i];
        hash ^= byte;
        hash *= FNV_PRIME;
    }
    return hash;
}

apfl_hash
apfl_hash_fnv1a(const void *data, size_t len)
{
    return apfl_hash_fnv1a_add(data, len, APFL_HASH_FNV1A_INIT);
}

#define HAS_CALLBACK(map, cb) ((map).callbacks.cb != NULL)
#define INVOKE_CALLBACK_NOARGS(map, cb) (map).callbacks.cb((map).callbacks.opaque)
#define INVOKE_CALLBACK(map, cb, ...) (map).callbacks.cb((map).callbacks.opaque, __VA_ARGS__)

static bool
keys_eq(const struct apfl_hashmap map, const void *a, const void *b)
{
    if (HAS_CALLBACK(map, keys_eq)) {
        return INVOKE_CALLBACK(map, keys_eq, a, b);
    } else {
        return memcmp(a, b, map.keysize) == 0;
    }
}

static apfl_hash
calc_hash(const struct apfl_hashmap map, const void *key)
{
    if (HAS_CALLBACK(map, calc_hash)) {
        return INVOKE_CALLBACK(map, calc_hash, key);
    } else {
        return apfl_hash_fnv1a(key, map.keysize);
    }
}

static void
destroy_key(const struct apfl_hashmap map, void *key)
{
    if (HAS_CALLBACK(map, destroy_key)) {
        INVOKE_CALLBACK(map, destroy_key, key);
    }
}

static void
destroy_value(const struct apfl_hashmap map, void *value)
{
    if (HAS_CALLBACK(map, destroy_value)) {
        INVOKE_CALLBACK(map, destroy_value, value);
    }
}

static void
copy_key(const struct apfl_hashmap map, void *dest, void *src)
{
    if (HAS_CALLBACK(map, copy_key)) {
        INVOKE_CALLBACK(map, copy_key, dest, src);
    } else {
        memcpy(dest, src, map.keysize);
    }
}

static void
copy_value(const struct apfl_hashmap map, void *dest, void *src)
{
    if (HAS_CALLBACK(map, copy_value)) {
        INVOKE_CALLBACK(map, copy_value, dest, src);
    } else {
        memcpy(dest, src, map.valsize);
    }
}

#define CAP_GROW 5
static_assert(CAP_GROW >= 1, "CAP_GROW must be at least 1");

static size_t
calc_new_cap(size_t old_cap)
{
    return old_cap + CAP_GROW;
}

#define KVADDR(base, elemsize, off) (((char*)(base)) + ((elemsize)*(off)))

static bool
find_key_in_bucket(const struct apfl_hashmap map, struct apfl_hashmap_bucket *bucket, const void *key, size_t *off)
{
    size_t keysize = map.keysize;

    for (size_t i = 0; i < bucket->len; i++) {
        if (keys_eq(map, key, KVADDR(bucket->keys, keysize, i))) {
            *off = i;
            return true;
        }
    }

    return false;
}

void
apfl_hashmap_set_prepared(struct apfl_hashmap_prepared_set prep, void *value)
{
    size_t keysize = prep.map->keysize;
    size_t valsize = prep.map->valsize;

    if (prep.i < prep.bucket->len) {
        // Replace
        void *dest = KVADDR(prep.bucket->values, valsize, prep.i);
        destroy_value(*prep.map, dest);

        copy_value(*prep.map, dest, value);
    } else {
        // Append
        assert(prep.i == prep.bucket->len);
        assert(prep.i <= prep.bucket->keys_cap);
        assert(prep.i <= prep.bucket->values_cap);
        copy_key(*prep.map, KVADDR(prep.bucket->keys, keysize, prep.i), prep.key);
        copy_value(*prep.map, KVADDR(prep.bucket->values, valsize, prep.i), value);

        prep.bucket->len++;
    }
}

static bool
prepare_set_in_bucket(
    struct apfl_hashmap *map,
    struct apfl_hashmap_prepared_set *prep,
    struct apfl_hashmap_bucket *bucket,
    void *key
) {
    size_t keysize = map->keysize;
    size_t valsize = map->valsize;

    size_t i;
    if (find_key_in_bucket(*map, bucket, key, &i)) {
        *prep = (struct apfl_hashmap_prepared_set) {
            .map = map,
            .bucket = bucket,
            .key = key,
            .i = i,
        };
        return true;
    }

    if (bucket->keys_cap != bucket->values_cap) {
        assert(false); // A prepare set operation was called on the hashmap,
                       // after a previous prepare set operation failed.
        return false;
    }

    if (bucket->len <= bucket->keys_cap) {
        size_t new_cap = calc_new_cap(bucket->keys_cap);

        void *newmem;

        newmem = REALLOC_UNTYPED_ARRAY(
            map->allocator,
            bucket->keys,
            keysize,
            bucket->keys_cap,
            new_cap
        );
        if (newmem == NULL) {
            return false;
        }
        bucket->keys = newmem;
        bucket->keys_cap = new_cap;

        newmem = REALLOC_UNTYPED_ARRAY(
            map->allocator,
            bucket->values,
            valsize,
            bucket->values_cap,
            new_cap
        );
        if (newmem == NULL) {
            return false;
        }
        bucket->values = newmem;
        bucket->values_cap = new_cap;
    }

    *prep = (struct apfl_hashmap_prepared_set) {
        .map = map,
        .bucket = bucket,
        .key = key,
        .i = bucket->len,
    };
    return true;
}

static void *
peek_in_bucket(const struct apfl_hashmap map, struct apfl_hashmap_bucket *bucket, const void *key)
{
    size_t i;
    if (!find_key_in_bucket(map, bucket, key, &i)) {
        return NULL;
    }

    return KVADDR(bucket->values, map.valsize, i);
}

static bool
get_in_bucket(const struct apfl_hashmap map, struct apfl_hashmap_bucket *bucket, const void *key, void *value)
{
    void *value_ptr = peek_in_bucket(map, bucket, key);

    if (value_ptr == NULL) {
        return false;
    }

    if (value != NULL) {
        copy_value(map, value, value_ptr);
    }

    return true;
}

static struct apfl_hashmap_bucket *
bucket_by_key(struct apfl_hashmap map, const void *key)
{
    apfl_hash hash = calc_hash(map, key);
    return &map.buckets[hash % map.nbuckets];
}

static void
delete_in_bucket(struct apfl_hashmap map, struct apfl_hashmap_bucket *bucket, const void *key)
{
    size_t i;
    if (!find_key_in_bucket(map, bucket, key, &i)) {
        return;
    }

    size_t keysize = map.keysize;
    size_t valsize = map.valsize;

    destroy_key(map, KVADDR(bucket->keys, keysize, i));
    destroy_value(map, KVADDR(bucket->values, valsize, i));

    assert(bucket->len >= (i+1));
    memmove(
        KVADDR(bucket->keys, keysize, i),
        KVADDR(bucket->keys, keysize, i+1),
        (bucket->len - (i+1)) * keysize
    );
    memmove(
        KVADDR(bucket->values, valsize, i),
        KVADDR(bucket->values, valsize, i+1),
        (bucket->len - (i+1)) * valsize
    );

    assert(bucket->len > 0); // if len == 0, we would not have found an entry
    bucket->len--;
}

#define INITIAL_NBUCKETS 16

static bool
hashmap_init(
    struct apfl_hashmap *map,
    struct apfl_allocator allocator,
    struct apfl_hashmap_callbacks callbacks,
    size_t nbuckets,
    size_t keysize,
    size_t valsize
) {
    map->callbacks = callbacks;
    map->allocator = allocator;
    map->keysize = keysize;
    map->valsize = valsize;
    map->nbuckets = nbuckets;
    if ((map->buckets = ALLOC_LIST(allocator, struct apfl_hashmap_bucket, nbuckets)) == NULL) {
        return false;
    }

    for (size_t i = 0; i < nbuckets; i++) {
        map->buckets[i] = (struct apfl_hashmap_bucket) {
            .keys       = NULL,
            .values     = NULL,
            .len        = 0,
            .keys_cap   = 0,
            .values_cap = 0,
        };
    }

    return true;
}

bool
apfl_hashmap_init(
    struct apfl_hashmap *map,
    struct apfl_allocator allocator,
    struct apfl_hashmap_callbacks callbacks,
    size_t keysize,
    size_t valsize
) {
    return hashmap_init(map, allocator, callbacks, INITIAL_NBUCKETS, keysize, valsize);
}

void
apfl_hashmap_delete(struct apfl_hashmap *map, const void *key)
{
    delete_in_bucket(*map, bucket_by_key(*map, key), key);
}

void *
apfl_hashmap_peek(const struct apfl_hashmap *map, const void *key)
{
    return peek_in_bucket(*map, bucket_by_key(*map, key), key);
}

bool
apfl_hashmap_get(const struct apfl_hashmap *map, const void *key, void *value)
{
    return get_in_bucket(*map, bucket_by_key(*map, key), key, value);
}

bool
apfl_hashmap_set(struct apfl_hashmap *map, void *key, void *value)
{
    struct apfl_hashmap_prepared_set prep;
    if (!apfl_hashmap_prepare_set(map, &prep, key)) {
        return false;
    }
    apfl_hashmap_set_prepared(prep, value);
    return true;
}

bool
apfl_hashmap_prepare_set(
    struct apfl_hashmap *map,
    struct apfl_hashmap_prepared_set *prep,
    void *key
) {
    return prepare_set_in_bucket(map, prep, bucket_by_key(*map, key), key);
}

static void
destroy_bucket(const struct apfl_hashmap map, struct apfl_hashmap_bucket *bucket)
{
    for (size_t i = 0; i < bucket->len; i++) {
        destroy_key(map, KVADDR(bucket->keys, map.keysize, i));
        destroy_value(map, KVADDR(bucket->values, map.valsize, i));
    }
    FREE_UNTYPED_ARRAY(map.allocator, bucket->keys, map.keysize, bucket->keys_cap);
    FREE_UNTYPED_ARRAY(map.allocator, bucket->values, map.valsize, bucket->values_cap);
    bucket->len = 0;
    bucket->keys_cap = 0;
    bucket->values_cap = 0;
}

size_t
apfl_hashmap_count(const struct apfl_hashmap map)
{
    size_t count = 0;
    for (size_t i = 0; i < map.nbuckets; i++) {
        count += map.buckets[i].len;
    }
    return count;
}

void
apfl_hashmap_deinit(struct apfl_hashmap *map)
{
    if (map == NULL) {
        return;
    }

    if (map->buckets != NULL) {
        for (size_t i = 0; i < map->nbuckets; i++) {
            destroy_bucket(*map, &map->buckets[i]);
        }
        FREE_LIST(map->allocator, map->buckets, map->nbuckets);
        map->buckets = NULL;
    }
}

struct apfl_hashmap
apfl_hashmap_move(struct apfl_hashmap *src)
{
    struct apfl_hashmap out = *src;
    src->buckets = NULL;
    src->nbuckets = 0;
    return out;
}

bool
apfl_hashmap_copy(struct apfl_hashmap *dst, struct apfl_hashmap src)
{
    size_t keysize = src.keysize;
    size_t valsize = src.valsize;

    struct apfl_hashmap_callbacks dst_callbacks = src.callbacks;

    if (!hashmap_init(dst, src.allocator, dst_callbacks, src.nbuckets, keysize, valsize)) {
        return false;
    }

    for (size_t i = 0; i < dst->nbuckets; i++) {
        struct apfl_hashmap_bucket *srcbucket = &src.buckets[i];
        struct apfl_hashmap_bucket *dstbucket = &dst->buckets[i];
        size_t len = srcbucket->len;

        if (len == 0) {
            continue;
        }

        dstbucket->keys = ALLOC_UNTYPED_ARRAY(src.allocator, keysize, len);
        dstbucket->values = ALLOC_UNTYPED_ARRAY(src.allocator, valsize, len);
        if (dstbucket->keys == NULL || dstbucket->values == NULL) {
            FREE_UNTYPED_ARRAY(src.allocator, dstbucket->keys, keysize, len);
            dstbucket->keys = NULL;
            FREE_UNTYPED_ARRAY(src.allocator, dstbucket->values, valsize, len);
            dstbucket->values = NULL;

            goto fail;
        }

        dstbucket->keys_cap = len;
        dstbucket->values_cap = len;

        for (size_t j = 0; j < len; j++) {
            copy_key(
                *dst,
                KVADDR(dstbucket->keys, keysize, j),
                KVADDR(srcbucket->keys, keysize, j)
            );

            copy_value(
                *dst,
                KVADDR(dstbucket->values, valsize, j),
                KVADDR(srcbucket->values, valsize, j)
            );

            dstbucket->len++;
        }
    }

    return dst;

fail:
    apfl_hashmap_deinit(dst);
    return false;
}


static void
cursor_skip_empty_buckets(struct apfl_hashmap_cursor *cur)
{
    struct apfl_hashmap *map = cur->map;
    while (cur->bucket < map->nbuckets && map->buckets[cur->bucket].len == 0) {
        cur->bucket++;
    }
}

struct apfl_hashmap_cursor
apfl_hashmap_get_cursor(struct apfl_hashmap *map)
{
    struct apfl_hashmap_cursor cursor = {
        .map = map,
        .i = 0,
        .bucket = 0,
    };
    cursor_skip_empty_buckets(&cursor);
    return cursor;
}

bool
apfl_hashmap_cursor_is_end(struct apfl_hashmap_cursor cursor)
{
    return cursor.bucket >= cursor.map->nbuckets;
}

static struct apfl_hashmap_bucket *
cursor_get_bucket(struct apfl_hashmap_cursor cursor)
{
    return apfl_hashmap_cursor_is_end(cursor)
        ? NULL
        : &cursor.map->buckets[cursor.bucket];
}

void
apfl_hashmap_cursor_next(struct apfl_hashmap_cursor *cursor)
{
    struct apfl_hashmap_bucket *bucket = cursor_get_bucket(*cursor);

    if (bucket == NULL) {
        return; // End already reached
    }

    cursor->i++;
    if (cursor->i < bucket->len) {
        return;
    }

    cursor->bucket++;
    cursor->i = 0;
    cursor_skip_empty_buckets(cursor);
}

#define CURSOR_PEEK(cursor, out, bucketmemb, sizememb)              \
    struct apfl_hashmap_bucket *bucket = cursor_get_bucket(cursor); \
                                                                    \
    if (bucket == NULL) {                                           \
        return NULL; /* End already reached */                      \
    }                                                               \
                                                                    \
    if (cursor.i >= bucket->len) {                                  \
        return NULL;                                                \
    }                                                               \
                                                                    \
    size_t size = cursor.map->sizememb;                             \
                                                                    \
    return KVADDR(bucket->bucketmemb, size, cursor.i);

void *
apfl_hashmap_cursor_peek_key(struct apfl_hashmap_cursor cursor)
{
    CURSOR_PEEK(cursor, *key_ptr, keys, keysize)
}

void *
apfl_hashmap_cursor_peek_value(struct apfl_hashmap_cursor cursor)
{
    CURSOR_PEEK(cursor, *value_ptr, values, valsize)
}

bool
apfl_hashmap_cursor_get_key(struct apfl_hashmap_cursor cursor, void *key)
{
    void *key_ptr = apfl_hashmap_cursor_peek_key(cursor);
    if (key_ptr == NULL) {
        return false;
    }
    copy_key(*cursor.map, key, key_ptr);
    return true;
}

bool
apfl_hashmap_cursor_get_value(struct apfl_hashmap_cursor cursor, void *value)
{
    void *value_ptr = apfl_hashmap_cursor_peek_value(cursor);
    if (value_ptr == NULL) {
        return false;
    }
    copy_value(*cursor.map, value, value_ptr);
    return true;
}