#include "paging.h"

uint32_t page_directory[1024] __attribute__((aligned(4096)));
uint32_t page_tables[128][1024] __attribute__((aligned(4096)));

void enable_paging() {
    asm(
        "mov %%eax, %%cr3\n"
        "mov %%cr0, %%eax\n"
        "or $0x80000001, %%eax\n"
        "mov %%eax, %%cr0\n"
    ::"a"(page_directory));
}

// TODO Add a function which can dynamically add and remove pages,
// that keeps track internally of what's already used?

// NOTE This function cannot cross 4MB (1024 page) boundaries
void map_pages(uintptr_t page, uintptr_t num_pages, uintptr_t phy_address, char access_flags) {
    uintptr_t access = access_flags & 0x7;

    // Each page is 4K, Each page table is 4M
    uintptr_t page_dir_entry = page >> 10;
    uintptr_t page_index = page & 0x3ff;
    if (page_index + num_pages > 1024) return;

    // Address must be 4K aligned
    if (phy_address & 0x3ff) return;

    for (int i = 0; i < num_pages; i++)
        page_tables[page_dir_entry][page_index + i] = phy_address + (i * 0x1000) | access_flags;
}

// Marks a 4MB page directory index unused, *without* freeing physical memory
void mark_pages_unused(uint16_t page_dir_idx) {
    if (page_dir_idx >= 1024) return;
    // Supervisor, R/W, Not Present
    page_directory[page_dir_idx] = 2;
}

// Bitmap of physical memory usage
uint8_t PhyMapArr[1024/8];

// Marks a 4MB page directory index unused, along with underlying physical memory
void free_pages(uint16_t page_dir_idx) {
    // Get first entry of page table
    uint32_t *page_table = (uint32_t*)(page_directory[page_dir_idx] != 2 ? (uintptr_t)page_directory[page_dir_idx] & 0xfffff000 : 0);
    if (!page_table) return;

    // Free physical memory
    uint32_t phy_address = page_table[0] & 0xfffff000;
    uint32_t phy_address_4mb = phy_address >> 22;
    uint32_t map_idx = phy_address_4mb >> 3;
    uint32_t bit_idx = phy_address_4mb & 7;
    // Clear bit
    PhyMapArr[map_idx] &= ~(1 << bit_idx);

    // Mark unused
    mark_pages_unused(page_dir_idx);
}

// Gets an unused 4MB page directory index. Return -1 if no free pages.
uint16_t get_unused_pages() {
    for (int i = 0; i < 1024; i++)
        if (page_directory[i] == 2) return i;
    return -1;
}

// Gets physical address of 4MB unused area of physical memory. This will have to be mapped to virtual memory space to be used.
uint32_t get_phy_mem() {
    for (int i = 0; i < sizeof(PhyMapArr); i++) {
        // Fully used
        if (PhyMapArr[i] == 0xff) continue;
        // At least one bit must be free
        uint8_t bits = PhyMapArr[i];
        for (int j = 0; j < 8; j++)
            if (!(bits & (1 << j))) return (i << 25) + (j << 22);
    }
    return -1;
}

// User, R/W, Present
#define USERACCESS (4|2|1)
// Supervisor, R/W, Present
#define SUPERACCESS (2|1)

// Returns a pointer to a newly mapped 4MB memory chunk
uintptr_t add_pages(char access_flags) {
    // Get free physical space
    uint32_t phy_mem = get_phy_mem();
    if (phy_mem == -1) return -1;

    // Get free virtual space
    uint16_t free_ent = get_unused_pages();
    if ((int16_t)free_ent == -1) return -1;
    uintptr_t page = free_ent << 10;

    // Mark page directory entry present, pointing to page table
    page_directory[free_ent] = ((uintptr_t)&page_tables[free_ent]) | USERACCESS;
    
    // Map pages to physical address
    map_pages(page, 1024, phy_mem, access_flags);

    // Mark physical memory used
    uint32_t phy_address_4mb = phy_mem >> 22;
    uint32_t map_idx = phy_address_4mb >> 3;
    uint32_t bit_idx = phy_address_4mb & 7;
    // Set bit
    PhyMapArr[map_idx] |= 1 << bit_idx;

    // Return virtual address
    return page << 12;
}

// Copies 4MB groups of pages
void move_pages(uint16_t src_page_dir_idx, uint16_t dest_page_dir_idx) {
    if (src_page_dir_idx >= 1024 || dest_page_dir_idx >= 1024) return;
    page_directory[dest_page_dir_idx] = page_directory[src_page_dir_idx];
}

void init_paging() {
    for (int i = 0; i < 1024; i++)
        // Supervisor, R/W, Not Present
        page_directory[i] = 2;

    // 1024 Pages = 1MB
    // First MB: Real Mode
    // TODO Make some of this not accessible to usermode?
    map_pages(0x00000000 >> 12, 256, 0x00000000, USERACCESS);
    // Next 3MB: Kernel
    map_pages(0x00100000 >> 12, 768, 0x00100000, SUPERACCESS);
    // Next 4MB: Kernel
    map_pages(0x00400000 >> 12, 1024, 0x00400000, SUPERACCESS);

    // We aren't using page directory permissions
    for (int i = 0; i < 4; i++)
        page_directory[i] = ((uintptr_t)&page_tables[i]) | USERACCESS;

    // Mark all physical memory free TODO Check memory size
    for (int i = 0; i < sizeof(PhyMapArr); i++)
        PhyMapArr[i] = 0;
    // Mark lowest 8MB as used
    PhyMapArr[0] = 3;

    enable_paging();
}