#include "collections.h"
#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define MIN(A, B) ((A) <= (B) ? (A) : (B))

#define MAX(A, B) ((A) >= (B) ? (A) : (B))

#define ALIGN(VAL, ALIGN)                                                      \
    ((VAL) % (ALIGN) != 0 ? (VAL) + ((ALIGN) - (VAL) % (ALIGN)) : (VAL))

void *array_push(void **data,
    size_t *capacity,
    size_t *count,
    void const *elem,
    size_t elem_size)
{
    if (!*data) {
        *capacity = 8;
        *data = malloc(*capacity * elem_size);
        *count = 0;
    } else if (*count + 1 > *capacity) {
        *capacity *= 2;
        *data = realloc(*data, *capacity * elem_size);
    }
    void *ptr = &((unsigned char *)*data)[*count * elem_size];
    if (elem) {
        memcpy(ptr, elem, elem_size);
    }
    *count += 1;
    return ptr;
}

void *array_insert_at(void **data,
    size_t *capacity,
    size_t *count,
    size_t idx,
    void const *elem,
    size_t elem_size)
{
    if (idx >= *count) {
        return array_push(data, capacity, count, elem, elem_size);
    }
    if (*count + 1 > *capacity) {
        array_push(data, capacity, count, NULL, elem_size);
    }
    void *src_ptr = &((unsigned char *)*data)[idx * elem_size];
    void *dest_ptr = &((unsigned char *)*data)[(idx + 1) * elem_size];
    memmove(dest_ptr, src_ptr, (*count - idx) * elem_size);
    *count += 1;
    if (elem) {
        memcpy(src_ptr, elem, elem_size);
    }
    return src_ptr;
}

#define small_count_max 3

void smallarray_construct(struct smallarray *a)
{
    *a = (struct smallarray) { 0 };
}

void smallarray_destroy(struct smallarray *a)
{
    if ((a->smalldata[0] & 1) == 0 && a->data) {
        free(a->data);
    }
}

static bool sa_is_big(struct smallarray const *a)
{
    return a->data && (a->smalldata[0] & 1) == 0;
}

static void sa_push_big(struct smallarray *a, void *value)
{
    if (!sa_is_big(a)) {
        size_t capacity = 8;
        void **ptr = malloc(capacity * sizeof(void *));
        size_t count = 0;

        for (size_t i = 0; i < small_count_max; ++i) {
            if ((a->smalldata[i] & 1) == 0)
                break;
            ptr[i] = (void *)(a->smalldata[i] & ~1ull);
            count += 1;
        }

        a->capacity = capacity;
        a->data = ptr;
        a->count = count;
    }

    array_push(
        (void **)&a->data, &a->capacity, &a->count, &value, sizeof(void *));
}

static void sa_push_small(struct smallarray *a, void *value)
{
    for (size_t i = 0; i < small_count_max; ++i) {
        if (a->smalldata[i] & 1)
            continue;
        a->smalldata[i] = (size_t)value | 1;
        return;
    }
}

void smallarray_push(struct smallarray *a, void *value)
{
    assert((size_t)value % 2 == 0 && "pointer must be 2 bytes aligned");

    if (sa_is_big(a) || smallarray_count(a) >= small_count_max) {
        sa_push_big(a, value);
    } else {
        sa_push_small(a, value);
    }
}

size_t smallarray_count(struct smallarray const *a)
{
    if (sa_is_big(a)) {
        return a->count;
    } else {
        size_t count = 0;
        for (size_t i = 0; i < small_count_max; ++i) {
            if ((a->smalldata[i] & 1) == 0)
                break;
            count += 1;
        }
        return count;
    }
}

void *smallarray_get(struct smallarray *a, size_t idx)
{
    if (sa_is_big(a)) {
        return a->data[idx];
    } else {
        if (idx >= small_count_max || (a->smalldata[idx] & 1) == 0)
            return NULL;
        return (void *)(a->smalldata[idx] & ~1ull);
    }
}

void const *smallarray_get_const(struct smallarray const *a, size_t idx)
{
    return smallarray_get((struct smallarray *)a, idx);
}

static uint64_t hash_key(char const *data)
{
    // djb2

    uint64_t hash = 5381;
    unsigned char c;

    while ((c = (unsigned char)*data++)) {
        hash = ((hash << 5) + hash) + c; // hash * 33 + c
    }

    return hash;
}

static uint64_t hash_key_sized(char const *data, size_t size)
{
    // djb2

    uint64_t hash = 5381;
    unsigned char c;

    while (size-- > 0 && (c = (unsigned char)*data++)) {
        hash = ((hash << 5) + hash) + c; // hash * 33 + c
    }

    return hash;
}

void hashmap_construct(struct hashmap *t)
{
    *t = (struct hashmap) { NULL, 0, 0, NULL, 0, 0 };
}

void hashmap_destroy(struct hashmap *t)
{
    if (t->buckets)
        free(t->buckets);
    if (t->keys)
        free(t->keys);
}

static struct hash_entry *find_entry(struct hashmap *m, uint64_t hash)
{
    for (size_t i = 0; i < m->buckets_count; ++i) {
        struct hash_bucket *bucket = &m->buckets[i];
        if (hash < bucket->first_hash)
            break;
        if (hash > bucket->last_hash)
            continue;
        for (size_t j = 0; j < bucket->count; ++j) {
            struct hash_entry *entry = &bucket->entries[j];
            if (entry->hash == hash)
                return entry;
        }
    }
    return NULL;
}

static struct hash_bucket *insert_bucket_at(struct hashmap *m, size_t idx)
{
    struct hash_bucket bucket = {
        .entries = { { 0 } },
        .count = 0,
        .first_hash = 0,
        .last_hash = 0,
    };

    return array_insert_at((void **)&m->buckets,
        &m->buckets_capacity,
        &m->buckets_count,
        idx,
        &bucket,
        sizeof(struct hash_bucket));
}

static struct hash_entry *add_entry_to_bucket(
    struct hash_bucket *b, uint64_t hash)
{
    b->count += 1;
    b->first_hash = MIN(b->first_hash, hash);
    b->last_hash = MAX(b->last_hash, hash);
    return &b->entries[b->count - 1];
}

static struct hash_entry *make_entry(struct hashmap *m, uint64_t hash)
{
    assert(m->buckets_count >= 1 && m->buckets[0].count != 0);

    if (hash < m->buckets[0].first_hash
        && m->buckets[0].count < bucket_capacity) {

        return add_entry_to_bucket(&m->buckets[0], hash);
    }
    for (size_t i = 0; i < m->buckets_count; ++i) {
        struct hash_bucket *curr = &m->buckets[i];

        struct hash_bucket *next
            = i + 1 < m->buckets_count ? &m->buckets[i + 1] : NULL;

        if (next && hash >= next->first_hash)
            continue;

        if (curr->count < bucket_capacity) {
            return add_entry_to_bucket(curr, hash);
        } else if (next && next->count < bucket_capacity) {
            return add_entry_to_bucket(next, hash);
        } else {
            struct hash_bucket *b = insert_bucket_at(m, i + 1);
            b->first_hash = UINT64_MAX;
            return add_entry_to_bucket(b, hash);
        }
    }
    assert(false);
}

static void insert_key(
    struct hashmap *m, char const *key, size_t key_size, uint64_t hash)
{
    key_size = key_size != 0 ? key_size : strlen(key);
    char *key_value = malloc(key_size + 1);
    strncpy(key_value, key, key_size);
    key_value[key_size] = '\0';

    struct hash_key_entry key_entry = {
        .hash = hash,
        .value = key_value,
    };

    array_push((void **)&m->keys,
        &m->keys_capacity,
        &m->keys_count,
        &key_entry,
        sizeof(key_entry));
}

static void hashmap_set_internal(struct hashmap *m,
    char const *key,
    size_t key_size,
    uint64_t hash,
    void *value)
{
    if (m->buckets_count == 0) {
        struct hash_bucket *bucket = insert_bucket_at(m, 0);
        bucket->entries[0] = (struct hash_entry) { hash, value };
        bucket->count += 1;
        bucket->first_hash = hash;
        bucket->last_hash = hash;

        insert_key(m, key, key_size, hash);
        return;
    }

    struct hash_entry *entry = find_entry(m, hash);
    if (!entry) {
        entry = make_entry(m, hash);
        entry->hash = hash;

        insert_key(m, key, key_size, hash);
    }
    entry->value = value;
    return;
}

void hashmap_set(struct hashmap *m, char const *key, void *value)
{
    uint64_t hash = hash_key(key);
    hashmap_set_internal(m, key, 0, hash, value);
}

void hashmap_set_sized(
    struct hashmap *m, char const *key, size_t key_size, void *value)
{
    uint64_t hash = hash_key_sized(key, key_size);
    hashmap_set_internal(m, key, key_size, hash, value);
}

bool hashmap_has(struct hashmap *m, char const *key)
{
    uint64_t hash = hash_key(key);
    return find_entry(m, hash) != NULL;
}

static void *hashmap_get_internal(struct hashmap *m, uint64_t hash)
{
    struct hash_entry *entry = find_entry(m, hash);
    return entry ? entry->value : NULL;
}
void *hashmap_get(struct hashmap *m, char const *key)
{
    uint64_t hash = hash_key(key);
    return hashmap_get_internal(m, hash);
}
void *hashmap_get_sized(struct hashmap *m, char const *key, size_t key_size)
{
    uint64_t hash = hash_key_sized(key, key_size);
    return hashmap_get_internal(m, hash);
}

void *hashmap_get_hash(struct hashmap *m, size_t hash)
{
    return hashmap_get_internal(m, hash);
}

void const *hashmap_get_hash_const(struct hashmap const *m, size_t hash)
{
    return hashmap_get_internal((struct hashmap *)m, hash);
}

void hashmap_keys(
    struct hashmap const *m, struct hash_key_entry const **keys, size_t *count)
{
    *keys = m->keys;
    *count = m->keys_count;
}

struct blockalloc_block {
    unsigned char *data;
    size_t size;
};

void blockalloc_construct(struct blockalloc *a)
{
    *a = (struct blockalloc) { NULL, 0, 0, 0 };
}

void blockalloc_destroy(struct blockalloc *a)
{
    if (a->blocks)
        free(a->blocks);
}

void *blockalloc_alloc(struct blockalloc *a, size_t size, size_t align)
{
    size_t p = ALIGN(a->p, align);
    if (!a->blocks || p + size > a->blocks[a->count - 1].size) {
        size_t block_size = MAX(blockalloc_default_block, size);

        struct blockalloc_block block = {
            .data = malloc(block_size),
            .size = block_size,
        };

        array_push((void **)&a->blocks,
            &a->capacity,
            &a->count,
            &block,
            sizeof(struct blockalloc_block));
        a->p = 0;
        p = 0;
    }
    void *ptr = &a->blocks[a->count - 1].data[p];
    a->p = p + size;
    return ptr;
}