#include "allocator.h"
#include "panic.h"
#include <stdlib.h>
#include <string.h>

size_t align_size(size_t size)
{
    const size_t word_size = sizeof(size_t);
    return (size + word_size) / word_size * word_size;
}

void gc_buffer_construct(GCBuffer* buffer, size_t capacity)
{
    void* data = malloc(capacity);
    ASSERT(buffer != NULL);
    *buffer = (GCBuffer) {
        .data = data,
        .index = 0,
        .capacity = capacity,
    };
}

void gc_buffer_destroy(GCBuffer* buffer) { free(buffer->data); }

void gc_buffer_reset(GCBuffer* buffer) { buffer->index = 0; }

void* gc_buffer_alloc(GCBuffer* buffer, size_t size)
{
    size_t aligned_size = align_size(size);
    if (buffer->capacity - buffer->index < aligned_size) {
        return NULL;
    }
    void* ptr = &((size_t*)buffer->data)[buffer->index];
    buffer->index += aligned_size;
    return ptr;
}

void gc_buffer_resize(GCBuffer* buffer, size_t minimum_capacity)
{
    size_t aligned_min_cap = align_size(minimum_capacity);
    if (buffer->capacity >= aligned_min_cap) {
        return;
    }
    while (buffer->capacity < minimum_capacity) {
        buffer->capacity *= 2;
    }
    void* new_buffer = realloc(buffer->data, buffer->capacity);
    ASSERT(new_buffer);
    buffer->data = new_buffer;
}

void gc_construct(GC* allocator, size_t start_capacity)
{
    GCBuffer buffer_a;
    gc_buffer_construct(&buffer_a, start_capacity);
    GCBuffer buffer_b;
    gc_buffer_construct(&buffer_b, start_capacity);
    *allocator = (GC) {
        .alloc = gc_alloc,
        .realloc = gc_realloc,
        .buffer_a = buffer_a,
        .buffer_b = buffer_b,
        .select = GCBufferA,
    };
}

bool gc_buffer_contains(GCBuffer* buffer, void* address)
{
    return (size_t)address >= (size_t)buffer->data
        && (size_t)address < (size_t)buffer->data + (size_t)buffer->index;
}

void gc_destroy(GC* allocator)
{
    gc_buffer_destroy(&allocator->buffer_a);
    gc_buffer_destroy(&allocator->buffer_b);
}

void* gc_alloc(GC* allocator, size_t size)
{
    size_t full_size = size + sizeof(size_t);
    GCBuffer* buffer = gc_selected_buffer(allocator);
    void* ptr = gc_buffer_alloc(buffer, full_size);
    if (ptr == NULL) {
        return NULL;
    }
    size_t* alloc_size = (size_t*)ptr;
    *alloc_size = size;
    return &((size_t*)ptr)[1];
}

void* gc_realloc(GC* allocator, void* data, size_t size)
{
    size_t old_size = ((size_t*)data)[-1];
    void* new_ptr = gc_alloc(allocator, size);
    if (new_ptr == NULL) {
        return NULL;
    }
    memcpy(new_ptr, (size_t*)data - sizeof(size_t), old_size + sizeof(size_t));
    size_t* size_ptr = (size_t*)new_ptr;
    *size_ptr = size;
    return &((size_t*)new_ptr)[1];
}

void gc_collect(GC* allocator, void* entries, size_t entries_size)
{
    gc_select_other(allocator);
    GCBuffer* source = gc_other_buffer(allocator);
    gc_relocate_branches_buffer(allocator, entries, entries_size);
    gc_buffer_reset(source);
}

void gc_relocate_branches(GC* allocator, void* allocation)
{
    size_t size = ((size_t*)allocation)[-1];
    gc_relocate_branches_buffer(allocator, allocation, size);
}

void gc_relocate_branches_buffer(
    GC* allocator, void* branches, size_t branches_size
)
{
    GCBuffer* source = gc_other_buffer(allocator);
    for (size_t i = 0; i < branches_size; ++i) {
        if (gc_buffer_contains(source, &((size_t*)branches)[i])) {
            gc_relocate_allocation(allocator, branches);
            gc_relocate_branches(allocator, branches);
        }
    }
}

void gc_relocate_allocation(GC* allocator, void* allocation)
{
    GCBuffer* dest = gc_selected_buffer(allocator);
    void* ptr = &((size_t*)allocation)[-1];
    size_t size = ((size_t*)ptr)[0];
    size_t full_size = size + sizeof(size_t);
    void* new_ptr = gc_buffer_alloc(dest, full_size);
    if (new_ptr == NULL) {
        gc_buffer_resize(dest, dest->capacity + full_size);
        new_ptr = gc_buffer_alloc(dest, full_size);
        ASSERT(new_ptr != NULL);
    }
    size_t* selected_size_ptr = &((size_t*)new_ptr)[0];
    *selected_size_ptr = size;
    memcpy(&((size_t*)new_ptr)[1], allocation, size);
}

GCBuffer* gc_selected_buffer(GC* allocator)
{
    switch (allocator->select) {
        case GCBufferA:
            return &allocator->buffer_a;
        case GCBufferB:
            return &allocator->buffer_b;
    }
}

GCBuffer* gc_other_buffer(GC* allocator)
{
    switch (allocator->select) {
        case GCBufferA:
            return &allocator->buffer_b;
        case GCBufferB:
            return &allocator->buffer_a;
    }
}

void gc_select_other(GC* allocator)
{
    switch (allocator->select) {
        case GCBufferA:
            allocator->select = GCBufferB;
            break;
        case GCBufferB:
            allocator->select = GCBufferA;
            break;
    }
}