#include "memory.h"
#include "stdint.h"
#include "print.h"
#define MEMORY_POOL_SIZE 0xFFFFF  // Define the size of the memory pool
static uint8_t memory_pool[MEMORY_POOL_SIZE]; // Static memory pool
static uint8_t* memory_pool_end = memory_pool; // Pointer to the end of allocated memory

// Block header to store size and free status
typedef struct block_header {
    size_t size;
    struct block_header* next; // For linking free blocks
    int is_free;
} block_header_t;

#define HEADER_SIZE sizeof(block_header_t)

static block_header_t* free_list = NULL; // Head of the free list

// Custom malloc implementation with block headers
void* malloc(size_t size) {
    size = (size + 3) & ~3;  // Align size to 4 bytes
    
    // First, check the free list for a suitable block
    block_header_t* current = free_list;
    block_header_t* prev = NULL;
    
    while (current) {
        if (current->is_free && current->size >= size) {
            current->is_free = 0;  // Mark as used
            return (void*)((uint8_t*)current + HEADER_SIZE); // Return memory after header
        }
        prev = current;
        current = current->next;
    }

    // No suitable block found in free list, allocate new memory from pool
    if ((memory_pool_end + size + HEADER_SIZE) > (memory_pool + MEMORY_POOL_SIZE)) {
        printf("Failed to allocate memory! Requested: %u bytes\n", size);
        return NULL;
    }

    block_header_t* header = (block_header_t*)memory_pool_end;
    header->size = size;
    header->is_free = 0;
    header->next = NULL;

    memory_pool_end += size + HEADER_SIZE;

    return (void*)((uint8_t*)header + HEADER_SIZE);
}

// Custom free implementation to add block back to the free list
void free(void* ptr) {
    if (!ptr) {
        return;
    }

    // Get the block header before the allocated memory
    block_header_t* header = (block_header_t*)((uint8_t*)ptr - HEADER_SIZE);
    header->is_free = 1;

    // Add the block to the free list
    header->next = free_list;
    free_list = header;
}


void list_used_blocks() {
    uint8_t* current_ptr = memory_pool;

    printf("Listing all used blocks:\n");

    // Traverse the memory pool and print each block's information
    while (current_ptr < memory_pool_end) {
        block_header_t* header = (block_header_t*)current_ptr;
        if (!header->is_free) {
            printf("Used block at %p | Size: %u bytes\n", (void*)(current_ptr + HEADER_SIZE), header->size);
        }
        current_ptr += HEADER_SIZE + header->size;
    }

    printf("End of used blocks.\n");
}




void* memcpy(void* dst, const void* src, uint16_t num) {
    uint8_t* d = (uint8_t*)dst;
    const uint8_t* s = (const uint8_t*)src;
    while (num--) {
        *d++ = *s++;
    }
    return dst;
}

void* memset(void* ptr, int value, uint16_t num) {
    uint8_t* p = (uint8_t*)ptr;
    while (num--) {
        *p++ = (uint8_t)value;
    }
    return ptr;
}

int memcmp(const void* ptr1, const void* ptr2, uint16_t num) {
    const uint8_t* p1 = (const uint8_t*)ptr1;
    const uint8_t* p2 = (const uint8_t*)ptr2;
    while (num--) {
        if (*p1 != *p2) {
            return *p1 - *p2;
        }
        p1++;
        p2++;
    }
    return 0;
}