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

uint32_t get_pte(void *virtualaddr) {
    uint32_t pdindex = (uint32_t)virtualaddr >> 22;
    uint32_t ptindex = (uint32_t)virtualaddr >> 12 & 0x3FF;
    uint32_t *pd = (uint32_t *)0xFFFFF000;

    if (!(pd[pdindex] & 0x1)) {
        printf("GET_PTE ERROR: PT not present\n");
        return 0;
    }

    uint32_t *pt = (uint32_t*)(0xFFC00000 + (pdindex << 12));

    return pt[ptindex];
}

void *get_physaddr(void *virtualaddr) {
    unsigned long pdindex = (unsigned long)virtualaddr >> 22;
    unsigned long ptindex = (unsigned long)virtualaddr >> 12 & 0x03FF;

    unsigned long *pd = (unsigned long *)0xFFFFF000;
    // Here you need to check whether the PD entry is present.

    uint32_t *pt = (uint32_t*)(0xFFC00000 + (pdindex << 12));
    // Here you need to check whether the PT entry is present.

    return (void *)((pt[ptindex] & ~0xFFF) + ((unsigned long)virtualaddr & 0xFFF));
}

void map_page(void *physaddr, void *virtualaddr, unsigned int flags) {
    // Make sure that both addresses are page-aligned.
    if ((unsigned long)physaddr & 0xFFF || (unsigned long)virtualaddr & 0xFFF) {
        // Error handling: addresses not page-aligned
        printf("MAP_PAGE ERROR: address not page-aligned\n");
        return;
    }


    unsigned long pdindex = (unsigned long)virtualaddr >> 22;
    unsigned long ptindex = (unsigned long)virtualaddr >> 12 & 0x03FF;

    unsigned long *pd = (unsigned long *)0xFFFFF000;
    // Here you need to check whether the PD entry is present.
    // When it is not present, you need to create a new empty PT and
    // adjust the PDE accordingly.

    if (!(pd[pdindex] & 0x1)) { // If PDE not present
        void *new_pt_phys = alloc_page();
        if (!new_pt_phys) {
            printf("failed to allocate a new page table\n");
            return;
        }

        unsigned int pde_flags = PAGE_PRESENT | PAGE_RW;
        if (flags & PAGE_USER) {
            pde_flags |= PAGE_USER;
        }
        pd[pdindex] = (unsigned long)new_pt_phys | pde_flags;

        uint32_t *pt_virt = (uint32_t*)(0xFFC00000 + (pdindex << 12));
        asm volatile("invlpg (%0)" : : "r" (pt_virt) : "memory");

        memset(pt_virt, 0, 4096);
    }

    uint32_t *pt = (uint32_t*)(0xFFC00000 + (pdindex << 12));
    // Here you need to check whether the PT entry is present.
    // When it is, then there is already a mapping present. What do you do now?

    if(pt[ptindex] & 0x1)
    {
        printf("this page 0x%X is already mapped\n", virtualaddr);
    }

    pt[ptindex] = ((unsigned long)physaddr) | (flags & 0xFFF); // Present

    // Now you need to flush the entry in the TLB
    // or you might not notice the change.
    asm volatile("invlpg (%0)" : : "r" (virtualaddr) : "memory");
}

void unmap_page(void *virtualaddr) {
    // Calculate page directory and page table indices
    unsigned long pdindex = (unsigned long)virtualaddr >> 22;
    unsigned long ptindex = (unsigned long)virtualaddr >> 12 & 0x03FF;

    // Get the page directory
    unsigned long *pd = (unsigned long *)0xFFFFF000;

    if (pd[pdindex] & 0x1) { // Check if the page directory entry is present
        // Get the page table
        unsigned long *pt = (unsigned long *)(0xFFC00000 + (pdindex << 12));

        // Clear the page table entry (remove the mapping)
        pt[ptindex] = 0;

        // Invalidate the TLB entry for this virtual address
        asm volatile("invlpg (%0)" : : "r" (virtualaddr) : "memory");
    }
}

/**
 * Maps a physical page to a virtual address in a SPECIFIED page directory.
 * pd_virt_addr: Kernel virtual address of the target page directory table.
 * phys_addr: Physical address of the page to map.
 * virt_addr: Virtual address to map to.
 * flags: PTE flags (e.g., PAGE_PRESENT | PAGE_USER | PAGE_RW).
 * Returns 0 on success, -1 on failure.
 */
int map_page_in_directory(uint32_t* pd_virt_addr, void* phys_addr, void* virt_addr, uint32_t flags) {
    uint32_t pde_idx = (uint32_t)virt_addr >> 22;
    uint32_t pte_idx = ((uint32_t)virt_addr >> 12) & 0x3FF;

    uint32_t pde = pd_virt_addr[pde_idx];
    uint32_t* pt_virt_for_modification; // Virtual address of the page table to write PTEs into

    if (!(pde & PAGE_PRESENT)) {
        // Page table not present, need to allocate one
        void* pt_phys_addr = alloc_page();
        if (!pt_phys_addr) {
            printf("map_page_in_directory: alloc_page failed for new page table (VA: %x)\n", (uint32_t)virt_addr);
            return -1;
        }

        // Get a kernel virtual address for this new page table to initialize it
        pt_virt_for_modification = (uint32_t*)phys_to_virt((uint32_t)pt_phys_addr);

        // Temporarily map this new PT page into the *current* address space to zero it.
        // This assumes map_page modifies the current CR3's address space.
        map_page(pt_phys_addr, (void*)pt_virt_for_modification, PAGE_PRESENT | PAGE_RW); 
        memset(pt_virt_for_modification, 0, PAGE_SIZE);
        // Now, point the PDE in the target page directory (pd_virt_addr) to the new page table.
        // The PDE stores the PHYSICAL address of the page table.
        pd_virt_addr[pde_idx] = (uint32_t)pt_phys_addr | PAGE_PRESENT | PAGE_RW | PAGE_USER; // Ensure PAGE_USER if PTEs will have USER bit
        // Note: If map_page for the PT itself isn't needed long-term in current PD, you might unmap it.
        // However, if phys_to_virt gives a stable kernel mapping, this explicit map/unmap might be simplified.
    } else {
        // Page table already exists. Get its physical address from PDE, then its kernel virtual address.
        pt_virt_for_modification = (uint32_t*)phys_to_virt(pde & ~0xFFF);
        // This assumes that any existing page table pointed to by a PDE is already 
        // accessible via its phys_to_virt address for modification.
    }

    // Set the PTE entry in this page table
    // Ensure phys_addr is page-aligned; flags should not include low 12 bits of phys_addr
    pt_virt_for_modification[pte_idx] = ((uint32_t)phys_addr & ~0xFFF) | (flags & 0xFFF);
    
    // When CR3 for the target task is loaded, TLB entries for virt_addr will be appropriately managed.
    // If virt_addr was previously mapped by another process using this same PD (unlikely for new tasks),
    // more complex TLB invalidation might be needed, but typically not for initial setup.
    // invlpg may be needed here if you are re-mapping an address in a PD already in use by another CPU or
    // if the PD is the current CR3 and the mapping changes. For setting up a new PD not yet in CR3, it's usually fine.

    return 0; // Success
}