#include "lib/print.h"
#include "multiboot.h"

#include <stddef.h>

extern uint8_t __kernel_start;
extern uint8_t __kernel_end;

extern uint32_t page_directory[1024];
extern void flush_tlb();
extern void disable_pae();

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

#define PAGE_TABLES  0xFFC00000
#define PAGE_DIRECTORY  0xFFFFF000

#define PAGE_PRESENT  0x1
#define PAGE_WRITABLE 0x2
#define PAGE_USER     0x4

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;
            }

            uint32_t kernel_phys_start = (uint32_t)&__kernel_start - 0xC0000000;
            uint32_t kernel_phys_end   = (uint32_t)&__kernel_end - 0xC0000000;
            for (uint32_t addr = usable_start; addr < usable_end; addr += 4096) {
                if (addr <= kernel_phys_end && addr >= kernel_phys_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);
    }
}

static void* p_alloc_frame() {
    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;
}

static void p_free_frame(void* addr) {
    mark_free((uint32_t)addr);
}

static void map_page(uint32_t vaddr, uint32_t paddr, uint32_t flags) {
    uint32_t pde_entry = vaddr >> 22;
    uint32_t pte_entry = (vaddr >> 12) & 0x3FF;

    uint32_t* page_directory = (uint32_t*)PAGE_DIRECTORY;
    if (page_directory[pde_entry] & PAGE_PRESENT == 0) {
        uint32_t* new_table = (uint32_t*)p_alloc_frame();
        page_directory[pde_entry] = (uint32_t)new_table | flags;
    }
}

static void* p_alloc(uint32_t size_b, uint32_t vstart) {
    uint32_t pages = (size_b + PAGE_SIZE - 1) / PAGE_SIZE;
    
    for (uint32_t i = 0; i < pages; i++) {
        uint32_t vaddr = vstart + i * PAGE_SIZE;
        uint32_t paddr = (uint32_t)p_alloc_frame();
        
        map_page(vaddr, paddr, PAGE_PRESENT | PAGE_WRITABLE);
    }

    return (void*)vstart;
}

void init_pmm(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);
    }
}

void init_page_tables() {
    uint32_t* kernel_table = p_alloc_frame();
    for (int i = 0; i < 1024; i++) {
        kernel_table[i] = (i * PAGE_SIZE) | 0x3;
    }
    page_directory[768] = (uint32_t)(kernel_table) | 0x3;
    page_directory[0] = 0;
    disable_pae();
    flush_tlb();
}

void* init_kheap() {

}

uint32_t total_free_memory() {
    return total_available_memory;
}