#include "multiboot.h"

#include <stddef.h>

extern uint8_t __kernel_start;
extern uint8_t __kernel_end;

#define PAGE_SIZE 4096
#define MAX_RAM   0xFFFFFFFF // 4GB
#define BITMAP_SIZE ((MAX_RAM / PAGE_SIZE) / 8)

static uint32_t max_address = 0;
static uint32_t total_available_memory = 0;
static uint8_t memory_bitmap[BITMAP_SIZE];

static void mark_free(uint32_t addr) {
    uint32_t page_index = addr / PAGE_SIZE;
    uint32_t byte_index = page_index / 8;
    uint32_t bit_index = page_index % 8;
    memory_bitmap[byte_index] &= ~(1 << bit_index);
}

static void mark_used(uint32_t addr) {
    uint32_t page_index = addr / PAGE_SIZE;
    uint32_t byte_index = page_index / 8;
    uint32_t bit_index = page_index % 8;
    memory_bitmap[byte_index] |= (1 << bit_index);
}

static uint32_t align(uint32_t addr, uint32_t len, uint32_t* usable_start, uint32_t* usable_end) {
    uint32_t start = addr;
    uint32_t end = addr + len;
    *usable_start = (start + 4095) & ~0xFFF;
    *usable_end = end & ~0xFFF;
    return *usable_end - *usable_start;
}

static void parse_multiboot_mmap(void* base, uint32_t limit) {
    multiboot_mmap_entry* entry = (multiboot_mmap_entry*)base;
    void* end = (void*)((uint32_t)base + limit);
    total_available_memory = 0;
    while ((void*)entry < end) {
        uint32_t entry_size = entry->size + sizeof(entry->size);
        if (entry->type == 1) {
            uint32_t region_start = (uint32_t)(entry->addr);
            uint32_t region_end = (uint32_t)(entry->addr + entry->len);
            if (region_end > max_address) {
                max_address = region_end;
            }

            uint32_t usable_start, usable_end, usable_size;
            usable_size = align(region_start, entry->len, &usable_start, &usable_end);
            if (usable_size <=  4096) {
                continue;
            }

            for (uint32_t addr = usable_start; addr < usable_end; addr += 4096) {
                if (addr <= (uint32_t)&__kernel_end && addr >= (uint32_t)&__kernel_start) {
                    continue;
                }

                if (addr >= (uint32_t)entry && addr <= (uint32_t)entry + entry_size) {
                    continue;
                }

                if (addr <= 0x200000) {
                    continue;
                }

                total_available_memory += 4096;

                mark_free(addr);
            }
        }
        entry = (multiboot_mmap_entry*)((uint32_t)entry + entry_size);
    }
}

void init_mm(void* mmap_base, uint32_t mmap_limit) {
    for (int i = 0; i < BITMAP_SIZE; i++) {
        memory_bitmap[i] = 0xFF;
    }

    if (mmap_base != NULL && mmap_limit > 0) {
        parse_multiboot_mmap(mmap_base, mmap_limit);
    }
}

uint32_t total_free_memory() {
    return total_available_memory;
}

void* alloc_page() {
    for (uint32_t i = 0; i < BITMAP_SIZE; i++) {
        uint8_t b = memory_bitmap[i];
        if (b == 0xFF) {
            continue;
        }

        for (int bit = 0; bit < 8; bit++) {
            if ((b & (1 << bit)) == 0) {
                uint32_t page_addr = (i * 8 + bit) * PAGE_SIZE;
                mark_used(page_addr);
                return (void*)page_addr;
            }
        }
    }

    return 0;
}

void free_page(void* addr) {
    mark_free((uint32_t)addr);
}