#include <stdint.h>
#include <stddef.h>
#include <stdint.h>
#include <stddef.h>
#include <assert.h>
/**
 * General purpose allocator.
 * - Allocated pointers are aligned to 16 bytes.
 * - Backing buffer must be aligned to 16 bytes start
 * - 4Gb - 16 is max for the allocator backing buffer.
 * - Allocation of > 0 bytes returns a size aligned up to 16 bytes
 * - Allocation of 0 bytes returns a valid pointer with 0 usable space
 * The allocator maintains a doubly linked list of blocks,
 * when blocks are free'd, adjacent blocks are merged to keep
 * fragmentation to a minimum
 */

typedef struct BlockHeader {
    uint32_t next;            // offset from base
    uint32_t prev;            // offset from base
    uint32_t size_and_flags;  // size | free_bit
    // Padding to ensure we always return 16 byte aligned pointers.
    uint8_t pad[4];
} BlockHeader;

#define ALIGNMENT sizeof(BlockHeader)
#define ALIGN_UP(n) (((n) + (ALIGNMENT - 1)) & ~(ALIGNMENT - 1))
#define ALIGN_DOWN(n) ((n) & ~(ALIGNMENT - 1))

#define NULL_BLOCK (0xffffffff)
#define BLOCK_FREE ((uint32_t)1)
_Static_assert(ALIGNMENT == 16, "sizeof(BlockHeader) is not 16");

typedef struct {
    uint8_t* base;            // backing store
    uint32_t size;            // total size
    BlockHeader* heap_start;  // first block
    BlockHeader* last_alloc;  // next-fit pointer
} SmallAllocator;

/* --- Header helpers --- */

static inline uint32_t block_size(const BlockHeader* b)
{
    return b->size_and_flags & ~BLOCK_FREE;
}

static inline int block_is_free(const BlockHeader* b)
{
    return (b->size_and_flags & BLOCK_FREE) != 0;
}

static inline void block_set_free(BlockHeader* b)
{
    b->size_and_flags |= BLOCK_FREE;
}

static inline void block_set_used(BlockHeader* b)
{
    b->size_and_flags &= ~BLOCK_FREE;
}

static inline void block_set_size(BlockHeader* b, uint32_t size)
{
    uint32_t flags = b->size_and_flags & BLOCK_FREE;
    b->size_and_flags = size | flags;
}

static inline BlockHeader* block_from_offset(const SmallAllocator* a, uint32_t off)
{
    if (off == NULL_BLOCK) {
        return NULL;
    }

    return (BlockHeader*)((uintptr_t)a->base + off);
}

static inline uint32_t bh_offset(SmallAllocator* a, BlockHeader* bh)
{
    if (bh == NULL) {
        return NULL_BLOCK;
    }

    return (uint32_t)((uintptr_t)bh - (uintptr_t)a->base);
}

void small_allocator_init(SmallAllocator* a, void* backing_store, uint32_t size)
{
    // backing store must be aligned, otherwise  all alignment guarantees are lies
    assert(((uintptr_t)backing_store & (ALIGNMENT - 1)) == 0);
    assert(size >= sizeof(BlockHeader));
    assert(size < 0xfffffff0);

    a->base = backing_store;
    a->size = size;

    BlockHeader* first = backing_store;
    uint32_t usable = ALIGN_DOWN(size - sizeof(BlockHeader));
    block_set_free(first);
    block_set_size(first, usable);
    first->next = NULL_BLOCK;
    first->prev = NULL_BLOCK;

    a->heap_start = first;
    a->last_alloc = first;
}

static void maybe_split_block(SmallAllocator* a, BlockHeader* block, uint32_t size)
{
    // assumes size is already aligned
    uint32_t bsize = block_size(block);

    if (bsize < size + sizeof(BlockHeader) + ALIGNMENT) return;

    uint8_t* base = (uint8_t*)block;
    BlockHeader* new_block = (BlockHeader*)(base + sizeof(BlockHeader) + size);

    uint32_t new_size = ALIGN_UP(bsize - size - sizeof(BlockHeader));

    block_set_size(block, size);
    block_set_used(block);

    new_block->size_and_flags = new_size | BLOCK_FREE;
    new_block->next = block->next;
    new_block->prev = bh_offset(a, block);

    BlockHeader* next = block_from_offset(a, new_block->next);
    if (next) {
        next->prev = bh_offset(a, new_block);
    }

    block->next = bh_offset(a, new_block);
}

static void merge_with_next(SmallAllocator* a, BlockHeader* block)
{
    BlockHeader* next = block_from_offset(a, block->next);
    if (next == NULL || !block_is_free(next)) return;

    uint32_t new_size = ALIGN_UP(block_size(block) + sizeof(BlockHeader) + block_size(next));

    block_set_size(block, new_size);
    block->next = next->next;
    BlockHeader* nextnext = block_from_offset(a, next->next);
    if (nextnext) nextnext->prev = bh_offset(a, block);
}

void* small_alloc(SmallAllocator* a, uint32_t size)
{
    size = ALIGN_UP(size);

    BlockHeader* start = a->last_alloc ? a->last_alloc : a->heap_start;
    BlockHeader* curr = start;

    do {
        if (block_is_free(curr) && block_size(curr) >= size) {
            maybe_split_block(a, curr, size);
            block_set_used(curr);
            a->last_alloc = curr;
            return (uint8_t*)curr + sizeof(BlockHeader);
        }
        curr = block_from_offset(a, curr->next) ? block_from_offset(a, curr->next) : a->heap_start;
    } while (curr != start);

    return NULL;
}

/* Free */

void small_free(SmallAllocator* a, void* ptr)
{
    if (!ptr) {
        return;
    }

    BlockHeader* block = (BlockHeader*)((uint8_t*)ptr - sizeof(BlockHeader));

    assert((uint8_t*)ptr >= a->base && (uint8_t*)ptr < a->base + a->size);
    // double free detection
    assert(!block_is_free(block));

    block_set_free(block);
    merge_with_next(a, block);

    BlockHeader* prev = block_from_offset(a, block->prev);
    if (prev != NULL && block_is_free(prev)) {
        merge_with_next(a, prev);

        if (a->last_alloc == block) {
            a->last_alloc = prev;
        }
    }
}

#if 0
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

void debug_alloc(const SmallAllocator* a)
{
    puts("--debug start--");
    for (BlockHeader* b = a->heap_start; b; b = block_from_offset(a, b->next)) {
        printf("Block size %u free %d prev %u next %u \n", block_size(b), block_is_free(b), b->prev, b->next);
    }
    puts("--debug end--");
}

[[gnu::aligned(16)]] char buff[1023 * 1023];
int main(void)
{
    printf("bh size %zu\n", sizeof(BlockHeader));
#define SZ 64
    SmallAllocator a;
    char* ptrs[SZ];

    small_allocator_init(&a, buff, sizeof buff);
    debug_alloc(&a);
    ptrs[0] = small_alloc(&a, 32);
    ptrs[1] = small_alloc(&a, 80);
    ptrs[2] = small_alloc(&a, 32);
    small_free(&a, ptrs[1]);
    ptrs[3] = small_alloc(&a, 12);
    ptrs[4] = small_alloc(&a, 12);
    debug_alloc(&a);

    small_free(&a, ptrs[0]);
    small_free(&a, ptrs[2]);
    small_free(&a, ptrs[3]);
    small_free(&a, ptrs[4]);

    for (int i = 0; i < SZ; i++) {
        ptrs[i] = small_alloc(&a, rand() % 256 + 1);
    }

    for (int i = 0; i < SZ; i++) {
        // if (rand() % 2 == 0) {
        small_free(&a, ptrs[i]);
        // }
    }
    debug_alloc(&a);
}

#endif