#include "Paging.hpp"
#include <Memory/PageFrameAllocator.hpp>
#include <Common/Panic.hpp>
#include <Memory/HHDM.hpp>
#include <CppLib/Spinlock.hpp>

namespace Memory::VMM {
    Paging* g_paging = nullptr;

    // Protects user page table modifications from concurrent SMP access
    static kcp::Mutex pagingLock;

    void LockUserPaging()   { pagingLock.Acquire(); }
    void UnlockUserPaging() { pagingLock.Release(); }

    extern "C" uint64_t KernelStartSymbol;

    std::uint64_t GetPhysKernelAddress(std::uint64_t virtualAddress) {
        return Paging::GetPhysAddr(GetCR3(), (std::uint64_t)virtualAddress, true);
    }

    Paging::Paging() {
        PML4 = (PageTable*)SubHHDM((PageTable*)Memory::g_pfa->AllocateZeroed());
    }

    void Paging::Init(std::uint64_t kernelBaseVirt, std::uint64_t kernelSize, const montauk::boot::MemoryMap& memMap,
                      const montauk::boot::Framebuffer& framebuffer) {
        // Map kernel
        Kt::KernelLogStream(Kt::DEBUG, "VMM") << "Paging::Init called with kernelBaseVirt as 0x" << base::hex << kernelBaseVirt;

        for (std::uint64_t pageAddr = kernelBaseVirt; pageAddr < (kernelBaseVirt + kernelSize); pageAddr += 0x1000) {
            Map(GetPhysKernelAddress(pageAddr), pageAddr);
        }

        // Map HHDM: map physical memory contiguously from 0 up to the top of
        // real backing store, so the kernel can reach any RAM/firmware byte at
        // phys+hhdmBase.  Mapping from 0 (rather than per-region) also covers
        // unlisted low gaps -- BIOS ROM at 0xE0000, the EBDA, etc. -- that the
        // kernel still touches through the HHDM.
        //
        // The upper bound is computed over real backing store ONLY: RAM, ACPI
        // tables, modules, the framebuffer, and so on.  Reserved/Unknown
        // regions are deliberately EXCLUDED from the bound.  Firmware commonly
        // reports the 64-bit PCI MMIO aperture as a Reserved region hundreds of
        // GiB above RAM (its base scales with the CPU's physical-address width);
        // letting that inflate maxPhysAddr would make this loop allocate
        // gigabytes of 4 KiB page tables and exhaust the frame allocator (OOM).
        // Such MMIO is mapped on demand via MapMMIO with the right cache
        // attributes -- never through the HHDM -- so it is correct to leave it
        // out here.  Reserved regions that sit *below* the bound (e.g. low BIOS
        // areas) are still mapped by the contiguous fill.
        uint64_t maxPhysAddr = 0;
        for (size_t i = 0; i < memMap.count; i++) {
            const auto& entry = memMap.regions[i];
            if (entry.kind == montauk::boot::MemoryKind::Reserved
             || entry.kind == montauk::boot::MemoryKind::Unknown) {
                continue;
            }
            uint64_t entryEnd = entry.base + entry.length;
            if (entryEnd > maxPhysAddr) maxPhysAddr = entryEnd;
        }
        maxPhysAddr = (maxPhysAddr + 0xFFF) & ~0xFFFULL;

        for (uint64_t pageAddr = 0; pageAddr < maxPhysAddr; pageAddr += 0x1000) {
            Map(pageAddr, HHDM(pageAddr));
        }

        // The linear framebuffer can sit in a PCI BAR above the RAM bound we
        // just mapped (e.g. in the 32-bit MMIO hole), and it is not guaranteed
        // to appear as a memory-map region.  It MUST be reachable through the
        // HHDM before we switch to these page tables, because the very next log
        // line (and PAT setup, and framebuffer init) renders to it via phys+hhdmBase.
        // Map its pages explicitly (write-back for now; the framebuffer module
        // upgrades them to write-combining once PAT is reprogrammed).
        if (framebuffer.address != 0) {
            uint64_t fbPhys  = Memory::SubHHDM(framebuffer.address);
            uint64_t fbBytes = framebuffer.height * framebuffer.pitch;
            uint64_t fbPages = (fbBytes + 0xFFF) / 0x1000;
            fbPhys &= ~0xFFFULL;
            for (uint64_t p = 0; p < fbPages; p++) {
                uint64_t phys = fbPhys + p * 0x1000;
                Map(phys, HHDM(phys));
            }
        }

        LoadCR3(PML4);
        Kt::KernelLogStream(Kt::OK, "VMM") << "Switched CR3";
    }

    PageTable* Paging::HandleLevel(VirtualAddress virtualAddress, PageTable* table, const size_t level) {
        PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[virtualAddress.GetIndex(level)]);

        if (!entry->Present) {
            void* newPage = Memory::g_pfa->AllocateZeroed();
            if (newPage == nullptr) {
                Panic("OOM in Paging::HandleLevel (kernel page table allocation)", nullptr);
            }
            uint64_t downLevelAddr = Memory::SubHHDM((uint64_t)newPage);

            entry->Present = true;
            entry->Writable = true;
            entry->Address = downLevelAddr >> 12;

            return (PageTable*)downLevelAddr;
        } else {
            return (PageTable*)(entry->Address << 12);
        }
    }

    PageTable* Paging::HandleLevelUser(VirtualAddress virtualAddress, PageTable* table, const size_t level) {
        PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[virtualAddress.GetIndex(level)]);

        if (!entry->Present) {
            void* newPage = Memory::g_pfa->AllocateZeroed();
            if (newPage == nullptr) {
                Panic("OOM in Paging::HandleLevelUser (page table allocation)", nullptr);
            }
            uint64_t downLevelAddr = Memory::SubHHDM((uint64_t)newPage);

            entry->Present = true;
            entry->Writable = true;
            entry->Supervisor = 1;  // User-accessible
            entry->Address = downLevelAddr >> 12;

            return (PageTable*)downLevelAddr;
        } else {
            // Ensure User bit is set on existing entries in the user path
            entry->Supervisor = 1;
            return (PageTable*)(entry->Address << 12);
        }
    }

    void Paging::Map(std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            Panic("Value that isn't page-aligned passed as address to Paging::Map!", nullptr);
        }

        VirtualAddress virtualAddressObj(virtualAddress);

        auto PML3 = HandleLevel(virtualAddressObj, PML4, 4);
        auto PML2 = HandleLevel(virtualAddressObj, PML3, 3);
        auto PML1 = HandleLevel(virtualAddressObj, PML2, 2);

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&PML1->entries[virtualAddressObj.GetPageIndex()]);

        pageEntry->Present = true;
        pageEntry->Writable = true;

        pageEntry->Address = physicalAddress >> 12;
    }

    void Paging::MapWC(std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            Panic("Value that isn't page-aligned passed as address to Paging::MapWC!", nullptr);
        }

        VirtualAddress virtualAddressObj(virtualAddress);

        auto PML3 = HandleLevel(virtualAddressObj, PML4, 4);
        auto PML2 = HandleLevel(virtualAddressObj, PML3, 3);
        auto PML1 = HandleLevel(virtualAddressObj, PML2, 2);

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&PML1->entries[virtualAddressObj.GetPageIndex()]);

        pageEntry->Present = true;
        pageEntry->Writable = true;
        pageEntry->WriteThrough = true;   // PWT=1, PCD=0 → PAT entry 1 = WC

        pageEntry->Address = physicalAddress >> 12;
    }

    void Paging::MapMMIO(std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            Panic("Value that isn't page-aligned passed as address to Paging::MapMMIO!", nullptr);
        }

        VirtualAddress virtualAddressObj(virtualAddress);

        auto PML3 = HandleLevel(virtualAddressObj, PML4, 4);
        auto PML2 = HandleLevel(virtualAddressObj, PML3, 3);
        auto PML1 = HandleLevel(virtualAddressObj, PML2, 2);

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&PML1->entries[virtualAddressObj.GetPageIndex()]);

        pageEntry->Present = true;
        pageEntry->Writable = true;
        pageEntry->CacheDisabled = true;
        pageEntry->WriteThrough = true;

        pageEntry->Address = physicalAddress >> 12;
    }

    void Paging::MapUser(std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            Panic("Value that isn't page-aligned passed as address to Paging::MapUser!", nullptr);
        }

        VirtualAddress virtualAddressObj(virtualAddress);

        auto PML3 = HandleLevelUser(virtualAddressObj, PML4, 4);
        auto PML2 = HandleLevelUser(virtualAddressObj, PML3, 3);
        auto PML1 = HandleLevelUser(virtualAddressObj, PML2, 2);

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&PML1->entries[virtualAddressObj.GetPageIndex()]);

        pageEntry->Present = true;
        pageEntry->Writable = true;
        pageEntry->Supervisor = 1;  // User-accessible

        pageEntry->Address = physicalAddress >> 12;
    }

    std::uint64_t Paging::CreateUserPML4() {
        // Allocate a new PML4
        void* newPage = Memory::g_pfa->AllocateZeroed();
        uint64_t newPml4Phys = Memory::SubHHDM((uint64_t)newPage);
        PageTable* newPml4 = (PageTable*)newPage;  // HHDM virtual address

        // Copy kernel-half entries (256-511) from the global PML4
        PageTable* kernelPml4 = (PageTable*)Memory::HHDM((uint64_t)g_paging->PML4);
        for (int i = 256; i < 512; i++) {
            newPml4->entries[i] = kernelPml4->entries[i];
        }

        return newPml4Phys;
    }

    bool Paging::MapUserIn(std::uint64_t pml4Phys, std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        pagingLock.Acquire();
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            pagingLock.Release();
            Panic("Non-aligned address in Paging::MapUserIn!", nullptr);
        }

        VirtualAddress va(virtualAddress);

        // Walk/create page tables from the given PML4, setting User bit at each level.
        // Returns nullptr on OOM.
        auto walkLevel = [](PageTable* table, uint64_t index) -> PageTable* {
            PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[index]);
            if (!entry->Present) {
                void* newPage = Memory::g_pfa->AllocateZeroed();
                if (newPage == nullptr) return nullptr;
                uint64_t newPhys = Memory::SubHHDM((uint64_t)newPage);
                entry->Present = true;
                entry->Writable = true;
                entry->Supervisor = 1;  // User-accessible
                entry->Address = newPhys >> 12;
                return (PageTable*)newPhys;
            } else {
                entry->Supervisor = 1;
                return (PageTable*)(entry->Address << 12);
            }
        };

        PageTable* pml4 = (PageTable*)pml4Phys;
        auto pml3 = walkLevel(pml4, va.GetL4Index());
        if (!pml3) { pagingLock.Release(); return false; }
        auto pml2 = walkLevel(pml3, va.GetL3Index());
        if (!pml2) { pagingLock.Release(); return false; }
        auto pml1 = walkLevel(pml2, va.GetL2Index());
        if (!pml1) { pagingLock.Release(); return false; }

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&pml1->entries[va.GetPageIndex()]);
        pageEntry->Present = true;
        pageEntry->Writable = true;
        pageEntry->Supervisor = 1;
        pageEntry->Address = physicalAddress >> 12;
        pagingLock.Release();
        return true;
    }

    bool Paging::MapUserInWC(std::uint64_t pml4Phys, std::uint64_t physicalAddress, std::uint64_t virtualAddress) {
        pagingLock.Acquire();
        if (virtualAddress % 0x1000 != 0 || physicalAddress % 0x1000 != 0) {
            pagingLock.Release();
            Panic("Non-aligned address in Paging::MapUserInWC!", nullptr);
        }

        VirtualAddress va(virtualAddress);

        auto walkLevel = [](PageTable* table, uint64_t index) -> PageTable* {
            PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[index]);
            if (!entry->Present) {
                void* newPage = Memory::g_pfa->AllocateZeroed();
                if (newPage == nullptr) return nullptr;
                uint64_t newPhys = Memory::SubHHDM((uint64_t)newPage);
                entry->Present = true;
                entry->Writable = true;
                entry->Supervisor = 1;
                entry->Address = newPhys >> 12;
                return (PageTable*)newPhys;
            } else {
                entry->Supervisor = 1;
                return (PageTable*)(entry->Address << 12);
            }
        };

        PageTable* pml4 = (PageTable*)pml4Phys;
        auto pml3 = walkLevel(pml4, va.GetL4Index());
        if (!pml3) { pagingLock.Release(); return false; }
        auto pml2 = walkLevel(pml3, va.GetL3Index());
        if (!pml2) { pagingLock.Release(); return false; }
        auto pml1 = walkLevel(pml2, va.GetL2Index());
        if (!pml1) { pagingLock.Release(); return false; }

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&pml1->entries[va.GetPageIndex()]);
        pageEntry->Present = true;
        pageEntry->Writable = true;
        pageEntry->Supervisor = 1;
        pageEntry->WriteThrough = true;   // PWT=1, PCD=0 -> PAT entry 1 = WC
        pageEntry->Address = physicalAddress >> 12;
        pagingLock.Release();
        return true;
    }

    void Paging::UnmapUserIn(std::uint64_t pml4Phys, std::uint64_t virtualAddress) {
        pagingLock.Acquire();
        VirtualAddress va(virtualAddress);

        // Walk without allocating -- if any level is absent, nothing is mapped.
        auto walkRead = [](PageTable* table, uint64_t index) -> PageTable* {
            PageTableEntry* entry = (PageTableEntry*)Memory::HHDM(&table->entries[index]);
            if (!entry->Present) return nullptr;
            return (PageTable*)(entry->Address << 12);
        };

        PageTable* pml4 = (PageTable*)pml4Phys;
        auto pml3 = walkRead(pml4, va.GetL4Index());
        if (!pml3) { pagingLock.Release(); return; }
        auto pml2 = walkRead(pml3, va.GetL3Index());
        if (!pml2) { pagingLock.Release(); return; }
        auto pml1 = walkRead(pml2, va.GetL2Index());
        if (!pml1) { pagingLock.Release(); return; }

        PageTableEntry* pageEntry = (PageTableEntry*)Memory::HHDM(&pml1->entries[va.GetPageIndex()]);
        if (!pageEntry->Present) { pagingLock.Release(); return; }

        // Clear the entire 8-byte PTE
        *(uint64_t*)pageEntry = 0;

        // Invalidate TLB for this virtual address
        asm volatile("invlpg (%0)" :: "r"(virtualAddress) : "memory");
        pagingLock.Release();
    }

    void Paging::FreeUserHalf(std::uint64_t pml4Phys) {
        // PageTable* pointers store PHYSICAL addresses (same convention as MapUserIn/UnmapUserIn).
        // Each individual entry access goes through HHDM() to get a valid virtual pointer.
        PageTable* pml4 = (PageTable*)pml4Phys;

        // Walk user-half entries (0-255); kernel-half (256-511) is shared and must not be touched.
        for (int i4 = 0; i4 < 256; i4++) {
            PageTableEntry* pml4e = (PageTableEntry*)Memory::HHDM(&pml4->entries[i4]);
            if (!pml4e->Present) continue;

            uint64_t pdptPhys = (uint64_t)pml4e->Address << 12;
            PageTable* pdpt = (PageTable*)pdptPhys;

            for (int i3 = 0; i3 < 512; i3++) {
                PageTableEntry* pdpte = (PageTableEntry*)Memory::HHDM(&pdpt->entries[i3]);
                if (!pdpte->Present) continue;

                uint64_t pdPhys = (uint64_t)pdpte->Address << 12;
                PageTable* pd = (PageTable*)pdPhys;

                for (int i2 = 0; i2 < 512; i2++) {
                    PageTableEntry* pde = (PageTableEntry*)Memory::HHDM(&pd->entries[i2]);
                    if (!pde->Present) continue;

                    uint64_t ptPhys = (uint64_t)pde->Address << 12;
                    PageTable* pt = (PageTable*)ptPhys;

                    // Free all leaf physical pages
                    for (int i1 = 0; i1 < 512; i1++) {
                        PageTableEntry* pte = (PageTableEntry*)Memory::HHDM(&pt->entries[i1]);
                        if (!pte->Present) continue;

                        // Skip MMIO/WC pages (not PFA-managed)
                        if (pte->WriteThrough || pte->CacheDisabled) continue;

                        uint64_t pagePhys = (uint64_t)pte->Address << 12;
                        if (pagePhys != 0) {
                            Memory::g_pfa->Free((void*)Memory::HHDM(pagePhys));
                        }
                    }

                    // Free the PT page itself
                    Memory::g_pfa->Free((void*)Memory::HHDM(ptPhys));
                }

                // Free the PD page
                Memory::g_pfa->Free((void*)Memory::HHDM(pdPhys));
            }

            // Free the PDPT page
            Memory::g_pfa->Free((void*)Memory::HHDM(pdptPhys));
        }
    }

    bool Paging::IsUserRangeAccessible(std::uint64_t pml4Phys, std::uint64_t virtualAddress,
                                       std::uint64_t size, bool requireWrite) {
        static constexpr uint64_t USER_SPACE_END = 0x0000800000000000ULL;

        if (size == 0) return true;
        if (virtualAddress >= USER_SPACE_END) return false;
        if (virtualAddress + size < virtualAddress) return false;
        if (virtualAddress + size > USER_SPACE_END) return false;

        auto pageAccessible = [&](uint64_t pageVa) -> bool {
            VirtualAddress va(pageVa);
            PageTable* pml4 = (PageTable*)pml4Phys;

            PageTableEntry* pml4e = (PageTableEntry*)Memory::HHDM(&pml4->entries[va.GetL4Index()]);
            if (!pml4e->Present || !pml4e->Supervisor) return false;
            if (requireWrite && !pml4e->Writable) return false;

            PageTable* pml3 = (PageTable*)(pml4e->Address << 12);
            PageTableEntry* pml3e = (PageTableEntry*)Memory::HHDM(&pml3->entries[va.GetL3Index()]);
            if (!pml3e->Present || !pml3e->Supervisor) return false;
            if (requireWrite && !pml3e->Writable) return false;
            if (pml3e->LargerPages) return true;

            PageTable* pml2 = (PageTable*)(pml3e->Address << 12);
            PageTableEntry* pml2e = (PageTableEntry*)Memory::HHDM(&pml2->entries[va.GetL2Index()]);
            if (!pml2e->Present || !pml2e->Supervisor) return false;
            if (requireWrite && !pml2e->Writable) return false;
            if (pml2e->LargerPages) return true;

            PageTable* pml1 = (PageTable*)(pml2e->Address << 12);
            PageTableEntry* pte = (PageTableEntry*)Memory::HHDM(&pml1->entries[va.GetPageIndex()]);
            if (!pte->Present || !pte->Supervisor) return false;
            if (requireWrite && !pte->Writable) return false;
            return true;
        };

        uint64_t page = virtualAddress & ~0xFFFULL;
        uint64_t lastPage = (virtualAddress + size - 1) & ~0xFFFULL;
        for (;;) {
            if (!pageAccessible(page)) return false;
            if (page == lastPage) break;
            page += 0x1000;
        }

        return true;
    }

    // Mask to extract only the 40-bit physical address from a PTE Address field
    // (bits 12-51 of the PTE). The PageTableEntry::Address field is 52 bits wide
    // and includes NX, PK, and software-available bits that must be stripped.
    static constexpr uint64_t kPhysAddrMask = 0xFFFFFFFFFFULL; // 40 bits

    std::uint64_t Paging::GetPhysAddr(std::uint64_t pml4, std::uint64_t virtualAddress, bool use40BitL1) {
        VirtualAddress virtualAddressObj(virtualAddress);

        PageTable* pml4Virt = (PageTable*)HHDM(pml4);

        PageTableEntry* pml4_entry = &pml4Virt->entries[virtualAddressObj.GetL4Index()];

        PageTable* pml3 = (PageTable*)HHDM((pml4_entry->Address & kPhysAddrMask) << 12);
        PageTableEntry* pml3_entry = &pml3->entries[virtualAddressObj.GetL3Index()];

        // 1GB large page at PML3 level
        if (pml3_entry->LargerPages) {
            uint64_t physBase = ((uint64_t)(pml3_entry->Address & kPhysAddrMask) << 12) & ~0x3FFFFFFFULL;
            return physBase | (virtualAddress & 0x3FFFFFFFULL);
        }

        PageTable* pml2 = (PageTable*)HHDM((pml3_entry->Address & kPhysAddrMask) << 12);
        PageTableEntry* pml2_entry = &pml2->entries[virtualAddressObj.GetL2Index()];

        // 2MB large page at PML2 level
        if (pml2_entry->LargerPages) {
            uint64_t physBase = ((uint64_t)(pml2_entry->Address & kPhysAddrMask) << 12) & ~0x1FFFFFULL;
            return physBase | (virtualAddress & 0x1FFFFFULL);
        }

        PageTable* pml1 = (PageTable*)HHDM((pml2_entry->Address & kPhysAddrMask) << 12);

        if (use40BitL1 == true) {
            PageTableEntry40Bit* pml1_entry = (PageTableEntry40Bit*)&pml1->entries[virtualAddressObj.GetPageIndex()];
            return (uint64_t)pml1_entry->Address << 12;
        }

        PageTableEntry* pml1_entry = &pml1->entries[virtualAddressObj.GetPageIndex()];
        return (uint64_t)(pml1_entry->Address & kPhysAddrMask) << 12;
    }

    std::uint64_t Paging::GetPhysAddr(std::uint64_t virtualAddress) {
        return GetPhysAddr((std::uint64_t)PML4, virtualAddress, false);
    }

    void Paging::MapEfiRuntime(const montauk::boot::EfiInfo& efi) {
        if (efi.memoryMap == nullptr || efi.descriptorSize == 0) return;

        auto* base = (uint8_t*)efi.memoryMap;
        uint64_t descSize = efi.descriptorSize;
        uint64_t count = efi.memoryMapSize / descSize;

        struct EfiMemDesc {
            uint32_t Type;
            uint64_t PhysicalStart;
            uint64_t VirtualStart;
            uint64_t NumberOfPages;
            uint64_t Attribute;
        };

        static constexpr uint64_t EFI_MEMORY_RUNTIME = 0x8000000000000000ULL;

        for (uint64_t i = 0; i < count; i++) {
            auto* desc = (EfiMemDesc*)(base + i * descSize);
            if (desc->Attribute & EFI_MEMORY_RUNTIME) {
                for (uint64_t p = 0; p < desc->NumberOfPages; p++) {
                    uint64_t addr = desc->PhysicalStart + p * 0x1000;
                    Map(addr, addr);
                }
            }
        }

        FlushTLB();
    }
};