#include <assert.h>
#include <buddy.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

#define MAX_ORDER      22 // 2^22 * min_block_size = max allocatable block
#define MIN_BLOCK_SIZE 32 // Smallest allocatable block

typedef struct free_block {
    struct free_block *next;
} free_block_t;

typedef struct {
    uint8_t order;
} block_header_t;

static free_block_t *free_lists[MAX_ORDER + 1];

uintptr_t heap_start = (uintptr_t)&__heap_start;
uintptr_t heap_end = (uintptr_t)&__heap_end;

// Utility: round up to next power of two
static size_t next_pow2(size_t x) {
    if (x <= 1)
        return 1;
    x--;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
#if UINTPTR_MAX > 0xffffffff
    x |= x >> 32;
#endif
    return x + 1;
}

// Find order for size
static int size_to_order(size_t size) {
    size_t needed = (size < MIN_BLOCK_SIZE) ? MIN_BLOCK_SIZE : size;
    size_t block_size = next_pow2(needed);
    int    order = 0;
    while ((MIN_BLOCK_SIZE << order) < block_size) {
        order++;
    }
    return order;
}

// Initialize allocator
void buddy_init(uintptr_t start, uintptr_t end) {
    heap_start = (start + MIN_BLOCK_SIZE - 1) & ~(MIN_BLOCK_SIZE - 1);
    heap_end = end & ~(MIN_BLOCK_SIZE - 1);

    // Clear free lists
    for (int i = 0; i <= MAX_ORDER; i++) free_lists[i] = NULL;

    // Put the whole region as one big block
    size_t total_size = heap_end - heap_start;
    int    order = MAX_ORDER;
    while ((MIN_BLOCK_SIZE << order) > total_size && order > 0) order--;

    free_block_t *block = (free_block_t *)heap_start;
    block->next = NULL;
    free_lists[order] = block;
}

// Allocate memory
void *buddy_alloc(size_t size) {
    int order = size_to_order(size);

    for (int i = order; i <= MAX_ORDER; i++) {
        if (free_lists[i]) {
            // Remove block from this list
            free_block_t *block = free_lists[i];
            free_lists[i] = block->next;

            // Split down to requested order
            while (i > order) {
                i--;
                uintptr_t     buddy_addr = (uintptr_t)block + (MIN_BLOCK_SIZE << i);
                free_block_t *buddy = (free_block_t *)buddy_addr;
                buddy->next = free_lists[i];
                free_lists[i] = buddy;
            }

            /*
             * Here, we install a header into the block. This will allow us to use free, without
             * specifying size, since it can read the order from one byte before the block start.
             */
            {
                block_header_t *hdr = (block_header_t *)block;
                hdr->order = order;
                return (void *)(hdr + 1); // return pointer after header
                // return (void *)block; // Regular block return for reference
            }
        }
    }
    return NULL; // Out of memory
}

// Free memory
int buddy_free(void *ptr) {
    block_header_t *hdr = (block_header_t *)ptr - 1;
    int             order = hdr->order;

    assert_msg(order != MAX_ORDER, "The buddy freelist header seems to have been corrupted.");

    uintptr_t addr = (uintptr_t)ptr;

    while (order < MAX_ORDER) {
        uintptr_t buddy_addr = addr ^ (MIN_BLOCK_SIZE << order);

        // Check if buddy is in free list
        free_block_t **prev = &free_lists[order];
        free_block_t  *cur = free_lists[order];
        while (cur) {
            if ((uintptr_t)cur == buddy_addr) {
                // Remove buddy from list
                *prev = cur->next;
                // Merge
                if (buddy_addr < addr)
                    addr = buddy_addr;
                order++;
                goto try_merge;
            }
            prev = &cur->next;
            cur = cur->next;
        }
        break; // Buddy not free, stop
    try_merge:;
    }

    free_block_t *block = (free_block_t *)addr;
    block->next = free_lists[order];
    free_lists[order] = block;

    return order;
}