#include "lib/memory.h"
#include "memory/mm.h"

#define KHEAP_START 0xE0000000

typedef struct __attribute__((packed)) {
    size_t size;
    uint8_t free;
    void* next_blk;
} alloc_header_t;

void memcpy(const void* src, void* dst, size_t sz) {
    for (size_t i = 0; i < sz; i++) {
        ((uint8_t*)dst)[i] = ((uint8_t*)src)[i];
    }
}

void memset(const uint8_t* dst, uint8_t val, size_t sz) {
    for (size_t i = 0; i < sz; i++) {
        ((uint8_t*)dst)[i] = val;
    }
}

int strcmp(const char* s1, const char* s2) {
    while (*s1 && (*s1 == *s2)) {
        s1++;
        s2++;
    }
    return *(const uint8_t*)s1 - *(const uint8_t*)s2;
}

int strlen(const char* str) {
    int len = 0;
    while(*str != 0) {
        str++;
        len++;
    }
    return len;
}

int strncmp(const char *s1, const char *s2, size_t n) {
    // If n is 0, we don't compare anything and strings are "equal"
    while (n > 0) {
        // If characters don't match, or we hit the null-terminator of s1
        if (*s1 != *s2) {
            return *(unsigned char *)s1 - *(unsigned char *)s2;
        }
        
        // If we reached the end of the strings simultaneously
        if (*s1 == '\0') {
            break;
        }
        
        s1++;
        s2++;
        n--;
    }
    return 0;
}

char* strcpy(char* dest, const char* src) {
    char* saved_dest = dest;

    while ((*dest++ = *src++) != '\0') {}

    return saved_dest;
}

static void* kernel_heap_start = (void*)KHEAP_START;
static void* kernel_heap_end = (void*)KHEAP_START;
alloc_header_t* blk_list = NULL;

void kgrow(size_t pages) {
    for (size_t p = 0; p < pages; p++) {
        void* frame = p_alloc_frame();
        map_page((uint32_t)kernel_heap_end, (uint32_t)frame, PAGE_PRESENT | PAGE_WRITABLE);
        kernel_heap_end = (void*)((uint32_t)kernel_heap_end + PAGE_SIZE); 
    }
}

#define ALIGNMENT 4
#define ALIGN(size) (((size) + (ALIGNMENT - 1)) & ~(ALIGNMENT - 1))
void* kalloc(size_t sz) {
    sz = ALIGN(sz);

    if (kernel_heap_end == kernel_heap_start) {
        // probaably not great, but if we don't have a heap yet, allocate the first (or only) frame requested.
        // Then initialize the first block header for the allocator
        kgrow(4);
        
        blk_list = (alloc_header_t*)kernel_heap_start;
        blk_list->size = ((uint32_t)kernel_heap_end - (uint32_t)kernel_heap_start) - sizeof(alloc_header_t);
        blk_list->free = 1;
        blk_list->next_blk = NULL;
    }

    alloc_header_t* curr = blk_list;
    alloc_header_t* prev = NULL;

    while (curr != NULL) {
        if (curr->free && curr->size >= sz) {
            if (curr->size >= sz + sizeof(alloc_header_t) + ALIGNMENT) {
                alloc_header_t* new_blk = (alloc_header_t*)((uint32_t)curr + sizeof(alloc_header_t) + sz);
                
                new_blk->size = curr->size - sz - sizeof(alloc_header_t);
                new_blk->free = 1;
                new_blk->next_blk = curr->next_blk;
                
                curr->size = sz;
                curr->next_blk = new_blk;
            }
            
            curr->free = 0; 
            return (void*)((uint32_t)curr + sizeof(alloc_header_t));
        }
        prev = curr;
        curr = curr->next_blk;
    }

    // no more space, grow!!
    size_t required_space = sz + sizeof(alloc_header_t);
    size_t pages_to_grow = (required_space + PAGE_SIZE - 1) / PAGE_SIZE; 
    
    alloc_header_t* new_heap_blk = (alloc_header_t*)kernel_heap_end;
    
    kgrow(pages_to_grow);
    
    new_heap_blk->size = (pages_to_grow * PAGE_SIZE) - sizeof(alloc_header_t);
    new_heap_blk->free = 1;
    new_heap_blk->next_blk = NULL;
    
    if (prev != NULL) {
        prev->next_blk = new_heap_blk;
    }
    
    return kalloc(sz);
}

void kfree(void* ptr) {
    
}