// mm_virt.c - virtual memory manager
//
// Page tables:
//   One PML4 shared by all processes.
//   Upper half (bit 47 set) has U/S = 0 - ring-3 access always faults.
//   Lower half pages are PAGE_USER.
//
// Kernel pool:
//   A boundary-tag (header + footer) free-list allocator starting at POOL_BASE.
//   Pages are committed on demand from the PMM.
//
// User heap:
//   A single best-fit free list in the lower half.
//   Backed on demand from the PMM. No per-process heaps.
#include <stddef.h>
#include <string.h>
#include <mm_virt.h>
#include <ke_spinlock.h>

// ---------------------------------------------------------------------------
// Helpers - physical window access
// ---------------------------------------------------------------------------
static uint64_t gPhysBase = PHYS_MAP_BASE;

static inline void* PhysToVirt(PHYSICAL_ADDRESS p) {
    return (void *)(p + gPhysBase);
}
static inline uint64_t* PtePtr(PHYSICAL_ADDRESS tablePhys, unsigned idx) {
    return (uint64_t *)PhysToVirt(tablePhys) + idx;
}

// ---------------------------------------------------------------------------
// PML4 / page-table helpers
// ---------------------------------------------------------------------------
static PHYSICAL_ADDRESS gPml4Phys = 0;

// Extract the page-table index for each level from a virtual address.
#define PML4_IDX(va)  (((va) >> 39) & 0x1FF)
#define PDPT_IDX(va)  (((va) >> 30) & 0x1FF)
#define PDT_IDX(va)   (((va) >> 21) & 0x1FF)
#define PT_IDX(va)    (((va) >> 12) & 0x1FF)

// Ensure a page-table entry at 'parentPhys[index]' points to a child table.
// Allocates and zeros a new frame if the entry is not present.
// Returns the physical address of the child table, or 0 on OOM.
static PHYSICAL_ADDRESS EnsureTable(PHYSICAL_ADDRESS parentPhys,
                                    unsigned index, uint64_t extraFlags) {
    uint64_t* entry = PtePtr(parentPhys, index);
    if (*entry & PAGE_PRESENT) {
        *entry |= extraFlags;
        return* entry & ~0xFFFULL;
    }
    PHYSICAL_ADDRESS childPhys = MmAllocateFrame();
    if (!childPhys) return 0;
    memset(PhysToVirt(childPhys), 0, PAGE_SIZE);
    *entry = childPhys | PAGE_PRESENT | PAGE_WRITE | extraFlags;
    return childPhys;
}

// ---------------------------------------------------------------------------
// Public mapping
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS MmMapPage(uintptr_t virt, PHYSICAL_ADDRESS phys, uint64_t flags) {
    unsigned userFlag = (flags & PAGE_USER) ? PAGE_USER : 0;

    PHYSICAL_ADDRESS pdpt = EnsureTable(gPml4Phys, PML4_IDX(virt), userFlag);
    if (!pdpt) return STATUS_NO_MEMORY;
    PHYSICAL_ADDRESS pdt  = EnsureTable(pdpt, PDPT_IDX(virt), userFlag);
    if (!pdt)  return STATUS_NO_MEMORY;
    PHYSICAL_ADDRESS pt   = EnsureTable(pdt,  PDT_IDX(virt),  userFlag);
    if (!pt)   return STATUS_NO_MEMORY;

    uint64_t* pte = PtePtr(pt, PT_IDX(virt));
    *pte = phys | flags;
    __invlpg((void *)virt);
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS MmUnmapPage(uintptr_t virt) {
    uint64_t* e4 = PtePtr(gPml4Phys, PML4_IDX(virt));
    if (!(*e4 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
    uint64_t* e3 = PtePtr(*e4 & ~0xFFFULL, PDPT_IDX(virt));
    if (!(*e3 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
    uint64_t* e2 = PtePtr(*e3 & ~0xFFFULL, PDT_IDX(virt));
    if (!(*e2 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
    uint64_t* pte = PtePtr(*e2 & ~0xFFFULL, PT_IDX(virt));
    *pte = 0;
    __invlpg((void *)virt);
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS MmSetPageProtect(uintptr_t virt, size_t size, uint64_t flags) {
    if (!size) return STATUS_SUCCESS;

    // Round down to page boundary, round size up.
    uintptr_t base = virt & PAGE_MASK;
    uintptr_t end  = (virt + size + PAGE_SIZE - 1) & PAGE_MASK;

    for (uintptr_t va = base; va < end; va += PAGE_SIZE) {
        uint64_t* e4 = PtePtr(gPml4Phys,          PML4_IDX(va));
        if (!(*e4 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
        uint64_t* e3 = PtePtr(*e4 & ~0xFFFULL,    PDPT_IDX(va));
        if (!(*e3 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
        uint64_t* e2 = PtePtr(*e3 & ~0xFFFULL,    PDT_IDX(va));
        if (!(*e2 & PAGE_PRESENT)) return STATUS_NOT_FOUND;
        uint64_t* pte = PtePtr(*e2 & ~0xFFFULL,   PT_IDX(va));
        if (!(*pte & PAGE_PRESENT)) return STATUS_NOT_FOUND;

        // Preserve the physical address, replace everything else.
        PHYSICAL_ADDRESS phys = *pte & ~0xFFFULL & ~PAGE_NX;
        *pte = phys | flags;
        __invlpg((void *)va);
    }

    return STATUS_SUCCESS;
}

__declspec(dllexport) PHYSICAL_ADDRESS MmGetPhysicalAddress(uintptr_t virt) {
    if (!gPml4Phys) return 0; // Not initialized

    uint64_t pml4_idx = (virt >> 39) & 0x1FF;
    uint64_t pdpt_idx = (virt >> 30) & 0x1FF;
    uint64_t pd_idx   = (virt >> 21) & 0x1FF;
    uint64_t pt_idx   = (virt >> 12) & 0x1FF;

    uint64_t* pml4 = PhysToVirt(gPml4Phys);
    if (!(pml4[pml4_idx] & PAGE_PRESENT)) return 0;

    uint64_t* pdpt = PhysToVirt(pml4[pml4_idx] & ~0xFFFULL);
    if (!(pdpt[pdpt_idx] & PAGE_PRESENT)) return 0;

    // Check for 1GB huge pages (bit 7 - PS)
    if (pdpt[pdpt_idx] & 0x80) {
        return (pdpt[pdpt_idx] & 0x000FFFFFC0000000ULL) + (virt & 0x3FFFFFFF);
    }

    uint64_t* pd = PhysToVirt(pdpt[pdpt_idx] & ~0xFFFULL);
    if (!(pd[pd_idx] & PAGE_PRESENT)) return 0;
    
    // Check for 2MB huge pages (bit 7 - PS)
    if (pd[pd_idx] & 0x80) {
        return (pd[pd_idx] & 0x000FFFFFFFE00000ULL) + (virt & 0x1FFFFF);
    }

    uint64_t* pt = PhysToVirt(pd[pd_idx] & ~0xFFFULL);
    if (!(pt[pt_idx] & PAGE_PRESENT)) return 0;

    // Extract the physical frame and add the offset within the page
    return (pt[pt_idx] & ~0xFFFULL) + (virt & 0xFFF);
}

// ---------------------------------------------------------------------------
// Boundary-tag allocator (used for both pool and user heap)
// ---------------------------------------------------------------------------
// Each allocation has a BLOCK_HDR immediately before the user data, and a
// BLOCK_FTR immediately after. Both carry the block size and a 'free' flag.
// This allows O(1) coalescing of adjacent free blocks.

#define BLOCK_MAGIC  0xB10CB10CUL

typedef struct BLOCK_HDR {
    uint32_t Magic;
    uint32_t Free;      // 1 = free, 0 = used
    size_t   Size;      // usable bytes (excludes header + footer)
    struct BLOCK_HDR* PrevFree;
    struct BLOCK_HDR* NextFree;
} BLOCK_HDR;

typedef struct BLOCK_FTR {
    size_t Size;
    uint32_t Free;
    uint32_t Magic;
} BLOCK_FTR;

#define HDR_SIZE  (sizeof(BLOCK_HDR))
#define FTR_SIZE  (sizeof(BLOCK_FTR))
#define OVERHEAD  (HDR_SIZE + FTR_SIZE)
#define MIN_BLOCK (OVERHEAD + sizeof(void*))

static inline BLOCK_FTR* HdrToFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR *)((char *)h + HDR_SIZE + h->Size);
}
static inline BLOCK_HDR* FtrToHdr(BLOCK_FTR* f) {
    return (BLOCK_HDR *)((char *)f - HDR_SIZE - f->Size);
}
static inline BLOCK_HDR* NextBlock(BLOCK_HDR* h) {
    return (BLOCK_HDR *)((char *)h + HDR_SIZE + h->Size + FTR_SIZE);
}
static inline BLOCK_FTR* PrevFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR *)((char *)h - FTR_SIZE);
}

typedef struct HEAP {
    BLOCK_HDR* FreeList;    // doubly-linked list of free blocks
    uintptr_t   Start;       // virtual base of the arena
    uintptr_t   Brk;         // current committed end (page-aligned)
    uintptr_t   Limit;       // hard upper bound
    uint64_t    MapFlags;    // PAGE_* flags for new committed pages
    SPINLOCK    Lock;
} HEAP;

// Extend heap by at least 'need' bytes (rounded up to page boundary).
static int HeapGrow(HEAP* h, size_t need) {
    size_t grow = (need + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (h->Brk + grow > h->Limit) return 0;

    for (uintptr_t va = h->Brk; va < h->Brk + grow; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmAllocateFrame();
        if (!frame) return 0;
        if (!KSUCCESS(MmMapPage(va, frame, h->MapFlags))) {
            MmFreeFrame(frame);
            return 0;
        }
        memset((void *)va, 0, PAGE_SIZE);
    }

    // Make the new pages a single free block.
    BLOCK_HDR* blk = (BLOCK_HDR *)h->Brk;
    blk->Magic    = BLOCK_MAGIC;
    blk->Free     = 1;
    blk->Size     = grow - OVERHEAD;
    blk->PrevFree = (BLOCK_HDR *)0;
    blk->NextFree = h->FreeList;
    if (h->FreeList) h->FreeList->PrevFree = blk;
    h->FreeList = blk;
    HdrToFtr(blk)->Magic = BLOCK_MAGIC;
    HdrToFtr(blk)->Free  = 1;
    HdrToFtr(blk)->Size  = blk->Size;

    h->Brk += grow;
    return 1;
}

// Unlink a free block from the free list.
static void FreeListRemove(HEAP* h, BLOCK_HDR* blk) {
    if (blk->PrevFree) blk->PrevFree->NextFree = blk->NextFree;
    else               h->FreeList = blk->NextFree;
    if (blk->NextFree) blk->NextFree->PrevFree = blk->PrevFree;
    blk->PrevFree = blk->NextFree = (BLOCK_HDR *)0;
}

static void HeapTryDecommitTail(HEAP* h, BLOCK_HDR* blk) {
    if (!blk || blk->Free == 0) return;

    while (h->Brk > h->Start) {
        uintptr_t blockEnd = (uintptr_t)blk + HDR_SIZE + blk->Size + FTR_SIZE;
        if (blockEnd != h->Brk) break;

        // Keep at least one page for the block header/footer and a small payload.
        if (blk->Size < PAGE_SIZE + sizeof(void*)) break;

        uintptr_t pageStart = h->Brk - PAGE_SIZE;
        if (pageStart <= (uintptr_t)blk) break; // would unmap the page containing the header

        PHYSICAL_ADDRESS phys = MmGetPhysicalAddress(pageStart);
        if (!phys) break;
        if (!KSUCCESS(MmUnmapPage(pageStart))) {
            break;
        }
        MmFreeFrame(phys);

        h->Brk   -= PAGE_SIZE;
        blk->Size -= PAGE_SIZE;
    }

    if (blk->Free) {
        HdrToFtr(blk)->Size  = blk->Size;
        HdrToFtr(blk)->Magic = BLOCK_MAGIC;
        HdrToFtr(blk)->Free  = 1;
    }
}

static void* HeapAlloc(HEAP* h, size_t size) {
    if (!size) return (void *)0;
    size = (size + 7) & ~7ULL;   // 8-byte align

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    // Best-fit search.
    BLOCK_HDR* best = (BLOCK_HDR *)0;
    for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
        if (cur->Size >= size) {
            if (!best || cur->Size < best->Size) best = cur;
            if (best->Size == size) break;
        }
    }

    if (!best) {
        if (!HeapGrow(h, size + OVERHEAD)) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void *)0;
        }
        // Re-search; HeapGrow inserted at the front.
        best = h->FreeList;
        while (best && best->Size < size) best = best->NextFree;
        if (!best) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void *)0;
        }
    }

    FreeListRemove(h, best);

    // Split the block if the remainder is big enough.
    if (best->Size >= size + OVERHEAD + sizeof(void*)) {
        BLOCK_HDR* rem = (BLOCK_HDR *)((char *)best + HDR_SIZE + size + FTR_SIZE);
        rem->Magic     = BLOCK_MAGIC;
        rem->Free      = 1;
        rem->Size      = best->Size - size - OVERHEAD;
        rem->PrevFree  = (BLOCK_HDR *)0;
        rem->NextFree  = h->FreeList;
        if (h->FreeList) h->FreeList->PrevFree = rem;
        h->FreeList    = rem;
        HdrToFtr(rem)->Magic = BLOCK_MAGIC;
        HdrToFtr(rem)->Free  = 1;
        HdrToFtr(rem)->Size  = rem->Size;

        best->Size = size;
    }

    best->Free = 0;
    HdrToFtr(best)->Free  = 0;
    HdrToFtr(best)->Magic = BLOCK_MAGIC;
    HdrToFtr(best)->Size  = best->Size;

    KeReleaseSpinlock(&h->Lock, flags);
    return (char *)best + HDR_SIZE;
}

static void HeapFree(HEAP* h, void* ptr) {
    if (!ptr) return;
    BLOCK_HDR* blk = (BLOCK_HDR *)((char *)ptr - HDR_SIZE);
    if (blk->Magic != BLOCK_MAGIC || blk->Free) return;   // double-free / corrupt

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    blk->Free = 1;

    // Coalesce with next block if it is free.
    BLOCK_HDR* next = NextBlock(blk);
    if ((uintptr_t)next < h->Brk && next->Magic == BLOCK_MAGIC && next->Free) {
        FreeListRemove(h, next);
        blk->Size += OVERHEAD + next->Size;
        HdrToFtr(blk)->Size  = blk->Size;
        HdrToFtr(blk)->Magic = BLOCK_MAGIC;
        HdrToFtr(blk)->Free  = 1;
    }

    // Coalesce with previous block if it is free.
    if ((uintptr_t)blk > h->Start) {
        BLOCK_FTR* prevFtr = PrevFtr(blk);
        if (prevFtr->Magic == BLOCK_MAGIC && prevFtr->Free) {
            BLOCK_HDR* prev = FtrToHdr(prevFtr);
            FreeListRemove(h, prev);
            prev->Size += OVERHEAD + blk->Size;
            HdrToFtr(prev)->Size  = prev->Size;
            HdrToFtr(prev)->Magic = BLOCK_MAGIC;
            HdrToFtr(prev)->Free  = 1;
            blk = prev;
        }
    }

    // Insert at head of free list.
    blk->PrevFree  = (BLOCK_HDR *)0;
    blk->NextFree  = h->FreeList;
    if (h->FreeList) h->FreeList->PrevFree = blk;
    h->FreeList = blk;

    HdrToFtr(blk)->Free  = 1;
    HdrToFtr(blk)->Magic = BLOCK_MAGIC;

    HeapTryDecommitTail(h, blk);

    KeReleaseSpinlock(&h->Lock, flags);
}

static void* HeapRealloc(HEAP* h, void* ptr, size_t new_size) {
    // If ptr is NULL, realloc acts like malloc
    if (!ptr) {
        return HeapAlloc(h, new_size);
    }

    // If size is 0, realloc acts like free
    if (new_size == 0) {
        HeapFree(h, ptr);
        return (void*)0;
    }

    // Retrieve the header of the existing block
    BLOCK_HDR* blk = (BLOCK_HDR *)((char *)ptr - HDR_SIZE);

    // Sanity checks to prevent operating on corrupted or already-freed pointers
    if (blk->Magic != BLOCK_MAGIC || blk->Free) {
        return (void*)0; 
    }

    // Align the requested size (matches HeapAlloc 8-byte alignment)
    size_t aligned_size = (new_size + 7) & ~7ULL;

    // If the current block is already big enough, we can reuse it in-place.
    // (Note: To save space, you could optionally split the block here if 
    // the new size is significantly smaller, similar to HeapAlloc's logic).
    if (blk->Size >= aligned_size) {
        return ptr;
    }

    // The block is too small. Allocate a new, larger block.
    void* new_ptr = HeapAlloc(h, aligned_size);
    if (!new_ptr) {
        return (void*)0; // Out of memory, original block is left untouched
    }

    // Copy the old data to the new block. 
    // We only copy up to the old block's usable size.
    memcpy(new_ptr, ptr, blk->Size);

    // Free the old block
    HeapFree(h, ptr);

    return new_ptr;
}

// ---------------------------------------------------------------------------
// Kernel pool
// ---------------------------------------------------------------------------
static HEAP gPoolHeap;

// ---------------------------------------------------------------------------
// User heap
// ---------------------------------------------------------------------------
static HEAP gUserHeap;

// ---------------------------------------------------------------------------
// Module init
// ---------------------------------------------------------------------------
void MmInitializeVirtual(PHYSICAL_ADDRESS bootPml4Phys, size_t poolSize) {
    gPml4Phys = bootPml4Phys;

    // Kernel pool heap.
    gPoolHeap.FreeList = (BLOCK_HDR *)0;
    gPoolHeap.Start    = POOL_BASE;
    gPoolHeap.Brk      = POOL_BASE;
    gPoolHeap.Limit    = POOL_LIMIT;
    gPoolHeap.MapFlags = PAGE_PRESENT | PAGE_WRITE | PAGE_NX;
    gPoolHeap.Lock     = (SPINLOCK)SPINLOCK_INIT;

    // Pre-commit the requested pool size.
    if (poolSize) {
        poolSize = (poolSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
        HeapGrow(&gPoolHeap, poolSize);
    }

    // User heap.
    gUserHeap.FreeList = (BLOCK_HDR *)0;
    gUserHeap.Start    = USER_BASE;
    gUserHeap.Brk      = USER_BASE;
    gUserHeap.Limit    = USER_LIMIT;
    gUserHeap.MapFlags = PAGE_PRESENT | PAGE_WRITE | PAGE_USER;
    gUserHeap.Lock     = (SPINLOCK)SPINLOCK_INIT;
}

void MmActivate(void) {
    __writecr3(gPml4Phys);
}

// ---------------------------------------------------------------------------
// Public pool API
// ---------------------------------------------------------------------------
__declspec(dllexport) void* MmAllocatePool(size_t size)                 { return HeapAlloc(&gPoolHeap, size); }
__declspec(dllexport) void  MmFreePool(void* ptr)                       { HeapFree(&gPoolHeap, ptr); }
__declspec(dllexport) void* MmReallocatePool(void* ptr, size_t newSize) { return HeapRealloc(&gPoolHeap, ptr, newSize); }

// ---------------------------------------------------------------------------
// Public user heap API
// ---------------------------------------------------------------------------
__declspec(dllexport) void* MmAllocateUserMemory(size_t size)                 { return HeapAlloc(&gUserHeap, size); }
__declspec(dllexport) void  MmFreeUserMemory(void* ptr, size_t size)          { (void)size; HeapFree(&gUserHeap, ptr); }
__declspec(dllexport) void* MmReallocateUserMemory(void* ptr, size_t newSize) { return HeapRealloc(&gUserHeap, ptr, newSize); }

size_t MmGetKernelPoolUsage(void) {
    // Returns bytes committed to the kernel pool
    return (size_t)(gPoolHeap.Brk - gPoolHeap.Start);
}

size_t MmGetUserHeapUsage(void) {
    // Returns bytes committed to the user heap
    return (size_t)(gUserHeap.Brk - gUserHeap.Start);
}