#include "../include/paging.h"
#include <stdint.h>
#include "../include/string.h"

uint32_t *kpage_dir;

uint32_t setup_tmp_pgdir(uint32_t magic, uint32_t info)
{
    int n, pd;
    uint32_t pagedir_address, memksize;
    uint32_t *page_table;
    struct multiboot_info *mbi;

    if(magic != MULTIBOOT_BOOTLOADER_MAGIC) {
        memksize = 4096;
    } else {
        mbi = (struct multiboot_info *)(PAGE_OFFSET + info);
        if(!(mbi->flags & MULTIBOOT_INFO_MEMORY)) {
            memksize = 4096;
        } else {
            memksize = (uint32_t)mbi->mem_upper + 1024;
            if(memksize > ((0xFFFFFFFF - PAGE_OFFSET) / 1024)) {
                memksize = (0xFFFFFFFF - PAGE_OFFSET) / 1024;
            }
        }
    }

    /* Address is PAGE_OFFSET plus memory size minus 4KB */
    pagedir_address = PAGE_OFFSET + (memksize * 1024) - memksize;
    pagedir_address = PAGE_ALIGN(pagedir_address);

    kpage_dir = (uint32_t *)pagedir_address;
    memset(kpage_dir, 0, PAGE_SIZE);

    pagedir_address += PAGE_SIZE;
    page_table = (uint32_t *)pagedir_address;
    memset(page_table, 0, memksize);

    for(n = 0; n < memksize / sizeof(uint32_t); n++) {
        page_table[n] = (n << PAGE_SHIFT) | PAGE_PRESENT | PAGE_RW;
        if(!(n % 1024)) {
            pd = n / 1024;
            kpage_dir[pd] = (unsigned int)(pagedir_address + (PAGE_SIZE * pd) + GDT_BASE) | PAGE_PRESENT | PAGE_RW;
            kpage_dir[GET_PGDIR(PAGE_OFFSET) + pd] = (unsigned int)(pagedir_address + (PAGE_SIZE * pd) + GDT_BASE) | PAGE_PRESENT | PAGE_RW;
        }
    }

    uint32_t pd_phys = (uint32_t)kpage_dir - PAGE_OFFSET;
    kpage_dir[1023] = pd_phys | PAGE_PRESENT | PAGE_RW; // Recursive mapping


    return (uint32_t)kpage_dir - PAGE_OFFSET;
}

void grub_memory_map(unsigned int magic, struct multiboot_info* mbi)
{
    /* Check bit 6 to see if we have a valid memory map */
    if(!(mbi->flags >> 6 & 0x1)) {
        printf("invalid memory map given by GRUB bootloader\n");
		while(1){}
    }

	/* Loop through the memory map and display the values */
    int i;
    for(i = 0; i < mbi->mmap_length; 
        i += sizeof(multiboot_memory_map_t)) 
    {
        multiboot_memory_map_t* mmmt = 
            (multiboot_memory_map_t*) (mbi->mmap_addr + i);

        printf("Start Addr: %x | Length: %x | Size: %x | Type: %x\n",
            mmmt->addr_low, mmmt->len_low, mmmt->size, mmmt->type);

        if(mmmt->type != MULTIBOOT_MEMORY_AVAILABLE) {
            /* 
             * Do something with this memory block!
             * BE WARNED that some of memory shown as availiable is actually 
             * actively being used by the kernel! You'll need to take that
             * into account before writing to memory!
             */
			printf("not available\n");
            uint32_t size_in_pages = mmmt->len_low / PAGE_SIZE;
            printf("size_in_pages: 0x%X\n", size_in_pages);

            uint32_t base_page = mmmt->addr_low / PAGE_SIZE;
            uint32_t page_count = mmmt->len_low / PAGE_SIZE;
            printf("page_base: %X, page_count: %X\n", base_page, page_count);
		    for (uint32_t j = 0; j < page_count; j++) {
                set_bit(base_page + j);
            }
        }
    }
    extern uint32_t _start, _end; // From linker script
    uint32_t kernel_start_phys = (uint32_t)&_start - 0xC0000000;
    uint32_t kernel_end_phys = (uint32_t)&_end - 0xC0000000;
    uint32_t kernel_start_page = kernel_start_phys / PAGE_SIZE;
    uint32_t kernel_end_page = (kernel_end_phys + PAGE_SIZE - 1) / PAGE_SIZE;
    printf("Kernel phys range: 0x%X to 0x%X, pages: 0x%X to 0x%X\n",
           kernel_start_phys, kernel_end_phys, kernel_start_page, kernel_end_page);

    for (uint32_t i = kernel_start_page; i < kernel_end_page; i++) {
        if (i < BITMAP_SIZE) {
            set_bit(i);
        } else {
            printf("Bitmap overflow at kernel page %u\n", i);
        }
    }
}

uint32_t* get_page_dir()
{
    uint32_t cr3_value;
    asm volatile("mov %%cr3, %0" : "=r" (cr3_value));
    return (uint32_t*) cr3_value;
}

void set_page_dir(uint32_t new_page_dir) {
    asm volatile("mov %0, %%cr3" : : "r" (new_page_dir));
}

uint32_t phys_to_virt(uint32_t phys)
{
    return phys+0xC0000000;
}

uint32_t* create_page_dir()
{
    void* ppp = alloc_page();
    memset(ppp, 0, 0x1000);

    uint32_t* current_pd = (uint32_t*)0xFFFFF000;  // Current PD (self-mapped)
    uint32_t* new_pd = (uint32_t*)phys_to_virt((uint32_t)ppp);

    for(int i = 768; i < 1023; i++)
    {
        if(current_pd[i] & PAGE_PRESENT) {
            new_pd[i] = current_pd[i];
        }
    }

    int vga_pde_index = 0; // Virtual address 0x000B8000 is covered by PDE[0]
    if (current_pd[vga_pde_index] & PAGE_PRESENT) {
        new_pd[vga_pde_index] = current_pd[vga_pde_index];
    }

    new_pd[1023] = (uint32_t)ppp | 0x03;  // Present + Read/Write

    return new_pd;
}

void destroy_page_dir(uint32_t* page_dir_virt) {
    if (page_dir_virt == NULL) {
        return;
    }

    // Iterate through user-space page directory entries (PDEs 0-767).
    // Kernel space (768+) is shared and should not be freed.
    for (int i = 0; i < 768; i++) {
        uint32_t pde = page_dir_virt[i];

        // Check if the page directory entry is present
        if (pde & PAGE_PRESENT) {
            // Get the physical address of the page table
            uint32_t pt_phys = pde & ~0xFFF;

            if(is_page_in_use(pt_phys / PAGE_SIZE)){
                // Convert it to a virtual address the kernel can access
                uint32_t* page_table_virt = (uint32_t*)phys_to_virt(pt_phys);

                // Iterate through all 1024 entries in this page table
                for (int j = 0; j < 1024; j++) {
                    uint32_t pte = page_table_virt[j];

                    // If the page table entry is present, free the physical page (frame) it points to
                    if (pte & PAGE_PRESENT) {
                        debug_log("freeing page %X\n", pte);
                        free_page((void*)(pte & ~0xFFF));
                    }
                }

                // After freeing all pages within the table, free the page table itself
                free_page((void*)pt_phys);
            }
        }
    }

    // Finally, free the page directory itself.
    // We need its physical address to pass to the physical memory manager.
    debug_log("freeing PD page %X\n", virt_to_phys(page_dir_virt));
    free_page((void*)virt_to_phys(page_dir_virt));
}