#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);
}

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

void hashmap_init(Hashmap* m)
{
    *m = (Hashmap) {
        .entries = nullptr,
        .entries_capacity = 8,
        .entries_size = 0,
    };
}

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

static KwPair* find_entry(Hashmap* m, const char* key)
{
    for (size_t i = 0; i < m->entries_size; ++i) {
        if (strcmp(m->entries[i].key, key) == 0) {
            return &m->entries[i];
        }
    }
    return nullptr;
}

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

    push_kwpair(m, &(KwPair) { strdup(key), val });
}

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

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];

    if (b->offset + size > b->cap) {
        size_t buffer_cap
            = a->default_bucket_cap > size ? a->default_bucket_cap : 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;

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

    return ptr;
}

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;
    void* new_buf = arena_alloc(a, size);
    memcpy(new_buf, p, 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);
}