#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

// Configuration
#define TOTAL_MEMORY   1024 // Total memory size in bytes
#define MIN_BLOCK_SIZE 16   // Smallest block size in bytes
#define MAX_LEVELS     10   //((int)(log2(TOTAL_MEMORY / MIN_BLOCK_SIZE)) + 1)

// Block structure to track memory blocks
typedef struct Block {
    bool          free; // Is the block free?
    size_t        size; // Size of the block
    struct Block *next; // Next block in the free list
} Block;

// Buddy allocator structure
typedef struct {
    char  *memory;                 // Memory pool
    Block *free_lists[MAX_LEVELS]; // Free lists for each level
    size_t total_size;             // Total memory size
} BuddyAllocator;

// Initialize the buddy allocator
void buddy_init(BuddyAllocator *allocator, size_t total_size) {
    allocator->total_size = total_size;
    allocator->memory = (char *)malloc(total_size);
    if (!allocator->memory) {
        printf("Memory allocation failed!\n");
        exit(1);
    }

    // Initialize free lists
    for (int i = 0; i < MAX_LEVELS; i++) {
        allocator->free_lists[i] = NULL;
    }

    // Create the initial block (entire memory)
    Block *initial_block = (Block *)allocator->memory;
    initial_block->free = true;
    initial_block->size = total_size;
    initial_block->next = NULL;

    // Add to the appropriate free list
    int level = (int)(log2(total_size / MIN_BLOCK_SIZE));
    allocator->free_lists[level] = initial_block;
}

// Get the level for a given size
int get_level(size_t size) {
    size_t adjusted_size = (size < MIN_BLOCK_SIZE) ? MIN_BLOCK_SIZE : size;
    return (int)ceil(log2((double)adjusted_size / MIN_BLOCK_SIZE));
}

// Allocate memory
void *buddy_alloc(BuddyAllocator *allocator, size_t size) {
    if (size == 0 || size > allocator->total_size) {
        return NULL;
    }

    // Adjust size to include block header and align to next power of 2
    size_t total_size = size + sizeof(Block);
    int    level = get_level(total_size);
    size_t block_size = MIN_BLOCK_SIZE << level;

    // Find a free block at the appropriate level
    int current_level = level;
    while (current_level < MAX_LEVELS && allocator->free_lists[current_level] == NULL) {
        current_level++;
    }

    if (current_level >= MAX_LEVELS) {
        return NULL; // No suitable block found
    }

    // Get the block
    Block *block = allocator->free_lists[current_level];
    allocator->free_lists[current_level] = block->next;

    // Split larger blocks if necessary
    while (current_level > level) {
        current_level--;
        block_size /= 2;

        // Create buddy block
        Block *buddy = (Block *)((char *)block + block_size);
        buddy->free = true;
        buddy->size = block_size;
        buddy->next = allocator->free_lists[current_level];
        allocator->free_lists[current_level] = buddy;

        block->size = block_size;
    }

    block->free = false;
    return (void *)((char *)block + sizeof(Block));
}

// Get the buddy of a block
Block *get_buddy(BuddyAllocator *allocator, Block *block) {
    size_t offset = (char *)block - allocator->memory;
    size_t buddy_offset = offset ^ block->size;
    return (Block *)(allocator->memory + buddy_offset);
}

// Merge free buddies
void merge_buddies(BuddyAllocator *allocator, Block *block, int level) {
    while (level < MAX_LEVELS - 1) {
        Block *buddy = get_buddy(allocator, block);
        // Check if buddy exists and is free
        Block *prev = NULL;
        Block *current = allocator->free_lists[level];
        while (current && current != buddy) {
            prev = current;
            current = current->next;
        }

        if (!current || !buddy->free) {
            break; // Buddy not free or not found
        }

        // Remove buddy from free list
        if (prev) {
            prev->next = buddy->next;
        } else {
            allocator->free_lists[level] = buddy->next;
        }

        // Merge into a larger block
        if ((char *)buddy < (char *)block) {
            block = buddy;
        }
        block->size *= 2;
        level++;
        block->next = allocator->free_lists[level];
        allocator->free_lists[level] = block;
    }
}

// Deallocate memory
void buddy_free(BuddyAllocator *allocator, void *ptr) {
    if (!ptr)
        return;

    // Get the block from the pointer
    Block *block = (Block *)((char *)ptr - sizeof(Block));
    block->free = true;

    // Add to free list and attempt to merge
    int level = (int)(log2(block->size / MIN_BLOCK_SIZE));
    block->next = allocator->free_lists[level];
    allocator->free_lists[level] = block;

    merge_buddies(allocator, block, level);
}

// Clean up the allocator
void buddy_destroy(BuddyAllocator *allocator) {
    free(allocator->memory);
    allocator->memory = NULL;
    for (int i = 0; i < MAX_LEVELS; i++) {
        allocator->free_lists[i] = NULL;
    }
}

// Example usage
int main() {
    BuddyAllocator allocator;
    buddy_init(&allocator, TOTAL_MEMORY);

    // Allocate memory
    void *ptr1 = buddy_alloc(&allocator, 100);
    printf("Allocated ptr1: %p\n", ptr1);

    void *ptr2 = buddy_alloc(&allocator, 50);
    printf("Allocated ptr2: %p\n", ptr2);

    // Free memory
    buddy_free(&allocator, ptr1);
    printf("Freed ptr1\n");

    buddy_free(&allocator, ptr2);
    printf("Freed ptr2\n");

    // Clean up
    buddy_destroy(&allocator);
    return 0;
}