#include "memory/mm.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 TEMP_MAPPING_VADDR 0xE0000000

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

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

// MODIFIES THE ACTIVE PAGE TABLE AT THE PHYSICAL ADDRESS LOADED INTO CR3 VIA RECURSIVE MAPPING
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* pd = (uint32_t*)PAGE_DIRECTORY;
    uint32_t* tablev = (uint32_t*)(PAGE_TABLES + (pde_entry * PAGE_SIZE));

    if ((pd[pde_entry] & PAGE_PRESENT) == 0) {
        uint32_t new_table = (uint32_t)p_alloc_frame();
        pd[pde_entry] = new_table | flags;
        flush_tlb();
        for (int i = 0; i < 1024; i++) {
            tablev[i] = 0;
        }
    }

    tablev[pte_entry] = paddr | flags;
    flush_tlb();
}

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_size = (uint32_t)&__kernel_end - (uint32_t)&__kernel_start;
    uint32_t num_pages = kernel_size / PAGE_SIZE + 1;
    uint32_t num_page_tables = (num_pages + 1023) / 1024;
    for (uint32_t table = 0; table < num_page_tables; table++) {
        uint32_t* kernel_table = p_alloc_frame();
        for (int i = 0; i < 1024; i++) {
            kernel_table[i] = (i * PAGE_SIZE) | 0x3;
        }
        page_directory[768 + table] = (uint32_t)(kernel_table) | 0x3;
    }
    page_directory[0] = 0;
    disable_pae();
    flush_tlb();
}

uint32_t total_free_memory() {
    return total_available_memory;
}

void* create_user_pd() {
    uint32_t pd_address = (uint32_t)p_alloc_frame();
    if (pd_address == 0) {
        return 0;
    }

    map_page(TEMP_MAPPING_VADDR, pd_address, PAGE_PRESENT | PAGE_WRITABLE);

    uint32_t* user_pd = (uint32_t*)TEMP_MAPPING_VADDR;
    uint32_t* kernel_pd = (uint32_t*)PAGE_DIRECTORY;

    for(int i = 0; i < 768; i++) {
        user_pd[i] = 0;
    }

    for(int i = 768; i < 1023; i++) {
        user_pd[i] = kernel_pd[i];
    }
    
    user_pd[1023] = (uint32_t)pd_address | PAGE_PRESENT | PAGE_WRITABLE;

    uint32_t pde = TEMP_MAPPING_VADDR >> 22;
    uint32_t pte = (TEMP_MAPPING_VADDR >> 12) & 0x3FF; 
    uint32_t* kernel_pt = (uint32_t*)(PAGE_TABLES + (pde * PAGE_SIZE));
    kernel_pt[pte] = 0;
    return (void*)pd_address;
}