#include "collection.h"
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static void array_push_generic(void** data,
    size_t* capacity,
    size_t* size,
    void* elem,
    size_t elem_size,
    void* alloc_fn(void*, size_t),
    void* realloc_fn(void*, void*, size_t),
    void* allocator)
{
    if (!*capacity) {
        *capacity = 4;
    }
    if (!*data) {
        *data = alloc_fn(allocator, *capacity * elem_size);
        *size = 0;
    }
    if (*size >= *capacity) {
        *capacity *= 2;
        *data = realloc_fn(allocator, *data, *capacity * elem_size);
    }
    memcpy((unsigned char*)*data + *size * elem_size, elem, elem_size);
    *size += 1;
}

static void* malloc_wrap(void* self, size_t size)
{
    (void)self;
    return malloc(size);
}

static void* realloc_wrap(void* self, void* p, size_t size)
{
    (void)self;
    return realloc(p, size);
}

void array_push(
    void** data, size_t* capacity, size_t* size, void* elem, size_t elem_size)
{
    array_push_generic(data,
        capacity,
        size,
        elem,
        elem_size,
        malloc_wrap,
        realloc_wrap,
        nullptr);
}

struct KwPair {
    char* key;
    unsigned char val[];
};

static KwPair* kwpair_at(Hashmap* m, size_t idx)
{
    return (KwPair*)((unsigned char*)m->entries + m->pair_size * idx);
}

static void push_kwpair(Hashmap* m, KwPair* pair)
{
    array_push((void**)&m->entries,
        &m->entries_cap,
        &m->entries_size,
        pair,
        m->pair_size);
}

void hashmap_init(Hashmap* m, size_t elem_size)
{
    *m = (Hashmap) {
        ARRAY_FIELD_INIT(entries),
        .elem_size = elem_size,
        .pair_size = sizeof(KwPair) + ((elem_size + 7) & ~7ul),
    };
}

void hashmap_deinit(Hashmap* m)
{
    if (m->entries) {
        for (size_t i = 0; i < m->entries_size; ++i) {
            free(kwpair_at(m, i)->key);
        }
        free(m->entries);
    }
}

static KwPair* find_entry(Hashmap* m, const char* key)
{
    for (size_t i = 0; i < m->entries_size; ++i) {
        KwPair* entry = kwpair_at(m, i);
        if (entry->key == key || strcmp(entry->key, key) == 0) {
            return entry;
        }
    }
    return nullptr;
}

void hashmap_set(Hashmap* m, const char* key, void* val)
{
    KwPair* existing = find_entry(m, key);
    if (existing) {
        memcpy(&existing->val, val, m->elem_size);
        return;
    }

    KwPair* entry = malloc(m->pair_size);
    entry->key = strdup(key);
    memcpy(&entry->val, val, m->elem_size);
    push_kwpair(m, entry);
    free(entry);
}

void* hashmap_get(Hashmap* m, const char* key)
{
    KwPair* existing = find_entry(m, key);
    return existing ? existing->val : nullptr;
}

void hashmap_for_each(
    Hashmap* m, void fn(void* self, const char* key, void* val), void* self)
{
    for (size_t i = 0; i < m->entries_size; ++i) {
        KwPair* entry = kwpair_at(m, i);
        fn(self, entry->key, &entry->val);
    }
}

void arena_init(Arena* a, size_t bucket_cap)
{
    *a = (Arena) {
        .buckets = nullptr,
        .buckets_cap = 0,
        .buckets_size = 0,
        .default_bucket_cap = bucket_cap,
    };

    array_push((void**)&a->buckets,
        &a->buckets_cap,
        &a->buckets_size,
        &(ArenaBucket) {
            .data = malloc(bucket_cap),
            .cap = bucket_cap,
            .offset = 0,
        },
        sizeof(ArenaBucket));
}

void arena_deinit(Arena* a)
{
    if (a->buckets) {
        for (size_t i = 0; i < a->buckets_size; ++i) {
            free(a->buckets[i].data);
        }
        free(a->buckets);
    }
}

void* arena_alloc(Arena* a, size_t size)
{
    ArenaBucket* b = &a->buckets[a->buckets_size - 1];

    size_t total_size = sizeof(size_t) + size;

    if (b->offset + total_size > b->cap) {
        size_t buffer_cap = a->default_bucket_cap > total_size
            ? a->default_bucket_cap
            : total_size;

        array_push((void**)&a->buckets,
            &a->buckets_cap,
            &a->buckets_size,
            &(ArenaBucket) {
                .data = malloc(buffer_cap),
                .cap = buffer_cap,
                .offset = 0,
            },
            sizeof(ArenaBucket));

        b = &a->buckets[a->buckets_size - 1];
    }

    unsigned char* ptr = b->data + b->offset;
    *(size_t*)ptr = size;

    b->offset += (total_size + 7ul) & ~7ul;

    return ptr + sizeof(size_t);
}

static void* arena_alloc_wrap(void* self, size_t size)
{
    Arena* a = self;
    return arena_alloc(a, size);
}

static void* arena_realloc_wrap(void* self, void* p, size_t size)
{
    Arena* a = self;
    size_t prev_size = *(size_t*)((unsigned char*)p - sizeof(size_t));
    void* new_buf = arena_alloc(a, size);
    memcpy(new_buf, p, prev_size);
    return new_buf;
}

void arena_array_push(Arena* arena,
    void** data,
    size_t* capacity,
    size_t* size,
    void* elem,
    size_t elem_size)
{
    array_push_generic(data,
        capacity,
        size,
        elem,
        elem_size,
        arena_alloc_wrap,
        arena_realloc_wrap,
        arena);
}

char* arena_strdup(Arena* a, const char* str)
{
    char* p = arena_alloc(a, strlen(str) + 1);
    strcpy(p, str);
    return p;
}

struct KeymapEntry {
    KeymapKey key;
    unsigned char value[];
};

void keymap_init(Keymap* m, size_t elem_size)
{
    size_t entry_size = sizeof(KeymapEntry) + (elem_size + 7ul & ~7ul);

    *m = (Keymap) {
        ARRAY_FIELD_INIT(entries),
        .elem_size = elem_size,
        .entry_size = entry_size,
        .temp_entry = malloc(entry_size),
    };
}

void keymap_deinit(Keymap* m)
{
    if (m->entries)
        free(m->entries);
    free(m->temp_entry);
}

static KeymapEntry* keymap_idx(Keymap* m, size_t idx)
{
    return (KeymapEntry*)((unsigned char*)m->entries + m->entry_size * idx);
}

struct SearchResult {
    KeymapEntry* entry;
    size_t idx;
};

static struct SearchResult keymap_search(
    Keymap* m, size_t begin, size_t end, KeymapKey key)
{
    if (begin == end)
        return (struct SearchResult) { nullptr, begin };

    size_t middle = (end - begin) / 2 + begin;
    KeymapEntry* entry = keymap_idx(m, middle);

    if (entry->key == key)
        return (struct SearchResult) { entry, middle };

    return key < entry->key ? keymap_search(m, begin, middle, key)
                            : keymap_search(m, middle + 1, end, key);
}

void keymap_set(Keymap* m, KeymapKey key, void* value)
{
    struct SearchResult result = keymap_search(m, 0, m->entries_size, key);

    if (result.entry) {
        memcpy(result.entry->value, value, m->elem_size);
        return;
    }

    printf("index = %ld\n", result.idx);

    if (m->entries_size) {
        if (keymap_idx(m, result.idx)->key < key) {
            printf("key is bigger\n");
        } else {
            printf("key is smaller\n");
        }
    }

    KeymapEntry* entry = m->temp_entry;
    memset(entry, 0, m->entry_size);

    array_push((void**)&m->entries,
        &m->entries_cap,
        &m->entries_size,
        entry,
        m->entry_size);

    entry->key = key;
    memcpy(&entry->value, value, m->elem_size);

    if (result.idx < m->entries_size - 1) {
        memmove(keymap_idx(m, result.idx + 1),
            keymap_idx(m, result.idx),
            (m->entries_size - result.idx - 1) * m->entry_size);
    }

    memcpy(keymap_idx(m, result.idx), entry, m->entry_size);

    for (size_t i = 0; i < m->entries_size; ++i) {
        KeymapEntry* entry = keymap_idx(m, i);
        printf("  %ld -> %d\n", entry->key, *(int*)entry->value);
    }
}

void* keymap_get(Keymap* m, KeymapKey key);