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

#include <stddef.h>
#include <lib/memory.h>

extern uint8_t __kernel_start;
extern uint8_t __kernel_end;

extern uint32_t page_directory[1024];

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) {
                entry = (multiboot_mmap_entry*)((uint32_t)entry + entry_size);
                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;
}

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*)CUR_PAGE_DIRECTORY;
    uint32_t* tablev = (uint32_t*)(CUR_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 unmap_page(uint32_t vaddr) {
    uint32_t pde_entry = vaddr >> 22;
    uint32_t pte_entry = (vaddr >> 12) & 0x3FF;
    uint32_t* tablev = (uint32_t*)(CUR_PAGE_TABLES + (pde_entry * PAGE_SIZE));
    tablev[pte_entry] = 0;
    flush_tlb();
}

int is_cow(uint32_t vaddr) {
    uint32_t pde_entry = vaddr >> 22;
    uint32_t pte_entry = (vaddr >> 12) & 0x3FF;
    uint32_t* pd = (uint32_t*)CUR_PAGE_DIRECTORY;
    uint32_t* tablev = (uint32_t*)(CUR_PAGE_TABLES + (pde_entry * PAGE_SIZE));

    if (!(pd[pde_entry] & PAGE_PRESENT)) {
        return 0;
    }

    if (tablev[pte_entry] & PAGE_WRITABLE) {
        return 0;
    }

    if (!(tablev[pte_entry] & PAGE_COW)) {
        return 0;
    }

    return 1;
}

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_new_pd() {
    uint32_t task_pd_addr = (uint32_t)p_alloc_frame();
    if (task_pd_addr == 0) {
        return NULL; // outta mem
    }

    map_page(TEMP_TASKPD_VADDR, task_pd_addr, PAGE_PRESENT | PAGE_WRITABLE);

    uint32_t* task_pd = (uint32_t*)TEMP_TASKPD_VADDR;
    uint32_t* parent_pd = (uint32_t*)CUR_PAGE_DIRECTORY;

    for(int i = 768; i < 1023; i++) {
        task_pd[i] = parent_pd[i];
    }

    task_pd[1023] = task_pd_addr | PAGE_PRESENT | PAGE_WRITABLE;

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

    unmap_page(TEMP_TASKPD_VADDR);

    return (void*)task_pd_addr;
}

void* clone_current_pd() {
    uint32_t task_pd_addr = (uint32_t)p_alloc_frame();
    if (task_pd_addr == 0) {
        return NULL; // outta mem
    }

    map_page(TEMP_TASKPD_VADDR, task_pd_addr, PAGE_PRESENT | PAGE_WRITABLE);

    uint32_t* task_pd = (uint32_t*)TEMP_TASKPD_VADDR;
    uint32_t* parent_pd = (uint32_t*)CUR_PAGE_DIRECTORY;

    for(int i = 768; i < 1023; i++) {
        task_pd[i] = parent_pd[i];
    }

    task_pd[1023] = task_pd_addr | PAGE_PRESENT | PAGE_WRITABLE;

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

    for (int i = 0; i < 768; i++) {
        uint32_t pde = parent_pd[i];
        if (!(pde & PAGE_PRESENT)) continue;

        uint32_t task_pt = (uint32_t)p_alloc_frame();
        task_pd[i] = task_pt | (pde & 0xFFF);

        uint32_t parent_pt = pde & 0xFFFFF000;
        
        map_page(TEMP_PT_DST, task_pt, PAGE_PRESENT | PAGE_WRITABLE);
        map_page(TEMP_PT_SRC, parent_pt, PAGE_PRESENT | PAGE_WRITABLE);

        uint32_t* pt_dst = (uint32_t*)TEMP_PT_DST;
        uint32_t* pt_src = (uint32_t*)TEMP_PT_SRC;

        for(int j = 0; j < 1024; j++) {
            uint32_t pte_src = pt_src[j];
            if (!(pte_src & PAGE_PRESENT)) continue;

            if (!(pte_src & PAGE_WRITABLE)) {
                pt_dst[j] = pt_src[j];
                continue;
            }

            // the src page is writable, mark the pte src as read only and COW, then copy to dst pte
            pt_src[j] &= ~(PAGE_WRITABLE);
            pt_src[j] |= PAGE_COW;
            flush_tlb();    // huge performance hit here, and probably everywhere else too where we aren't switching cr3s

            pt_dst[j] = pt_src[j];
        }

        unmap_page(TEMP_PT_DST);
        unmap_page(TEMP_PT_SRC);
    }

    unmap_page(TEMP_TASKPD_VADDR);

    return (void*)task_pd_addr;
}

void clear_current_pd() {
    
}