/*
 * memory.c – Kernel memory management.
 *
 * Implements three layers:
 *   PMM    – bitmap-based physical page-frame allocator backed by a
 *            16 MB pool obtained from the loader at boot.
 *   Paging – walks and creates 4-level x86-64 page tables; supports
 *            map, unmap, and virtual-to-physical translation.
 *   Heap   – first-fit free-list allocator with block splitting and
 *            bidirectional coalescing; 16-byte aligned.
 */

#include "memory.h"
#include "task.h"

/* Null-safe print helper used throughout the kernel. */
#define SAFE_PRINT(Boot, ...)                           \
    do {                                                \
        if ((Boot) != NULL && (Boot)->print != NULL) {  \
            (Boot)->print(__VA_ARGS__);                 \
        }                                               \
    } while (0)

/* ================================================================
 *  Physical Memory Manager  –  bitmap-based page-frame allocator
 * ================================================================ */

static UINT64  pmm_pool_base   = 0;
static UINTN   pmm_total_pages = 0;
static UINTN   pmm_free_count  = 0;
static UINT8   pmm_bitmap[PMM_POOL_PAGES / 8];
static BOOLEAN pmm_ready       = FALSE;

/* ================================================================
 *  PMM – bitmap helpers
 * ================================================================ */

/* Mark page `idx` as allocated. */
static void pmm_set_bit(UINTN idx)
{
    pmm_bitmap[idx / 8] |= (UINT8)(1U << (idx % 8));
}

/* Mark page `idx` as free. */
static void pmm_clear_bit(UINTN idx)
{
    pmm_bitmap[idx / 8] &= (UINT8)~(1U << (idx % 8));
}

/* Return TRUE if page `idx` is currently allocated. */
static BOOLEAN pmm_test_bit(UINTN idx)
{
    return (pmm_bitmap[idx / 8] & (1U << (idx % 8))) != 0;
}

/* ----------------------------------------------------------------
 *  PMM – public interface
 * ---------------------------------------------------------------- */

/*
 * Initialise the PMM: request PMM_POOL_PAGES from the loader via
 * BootInfo->alloc_pages() and set up the bitmap with all pages
 * marked free.
 */
void pmm_init(BootInfo *Boot)
{
    KSTATUS Status;
    UINT64 pool_addr = 0;
    UINTN i;

    /* Zero the bitmap – all pages start free */
    for (i = 0; i < sizeof(pmm_bitmap); i++) {
        pmm_bitmap[i] = 0;
    }

    if (Boot == NULL || Boot->alloc_pages == NULL) {
        SAFE_PRINT(Boot, L"PMM: page allocator unavailable\n\r");
        return;
    }

    Status = Boot->alloc_pages(PMM_POOL_PAGES, &pool_addr);
    if (Status != 0) {
        SAFE_PRINT(Boot, L"PMM: failed to allocate pool (%d pages), status=%ld\n\r",
                   (UINTN)PMM_POOL_PAGES, (UINT64)Status);
        return;
    }

    pmm_pool_base   = (UINT64)pool_addr;
    pmm_total_pages = PMM_POOL_PAGES;
    pmm_free_count  = PMM_POOL_PAGES;
    pmm_ready       = TRUE;

    SAFE_PRINT(Boot, L"  PMM : %d pages (%d KB) at 0x%lx\n\r",
               pmm_total_pages,
               (pmm_total_pages * PAGE_SIZE) / 1024,
               pmm_pool_base);
}

/* Allocate a single 4 KB page.  Returns physical address or 0. */
UINT64 pmm_alloc_page(void)
{
    UINTN i;

    if (!pmm_ready || pmm_free_count == 0) {
        return 0;
    }

    for (i = 0; i < pmm_total_pages; i++) {
        if (!pmm_test_bit(i)) {
            pmm_set_bit(i);
            pmm_free_count--;
            return pmm_pool_base + ((UINT64)i * PAGE_SIZE);
        }
    }

    return 0;
}

/* Free a single page previously returned by pmm_alloc_page(). */
void pmm_free_page(UINT64 phys_addr)
{
    UINTN idx;

    if (!pmm_ready) return;
    if (phys_addr < pmm_pool_base) return;

    idx = (UINTN)((phys_addr - pmm_pool_base) / PAGE_SIZE);
    if (idx >= pmm_total_pages) return;
    if (!pmm_test_bit(idx)) return;        /* already free */

    pmm_clear_bit(idx);
    pmm_free_count++;
}

/* Allocate `count` physically contiguous pages (first-fit). */
UINT64 pmm_alloc_pages(UINTN count)
{
    UINTN i, j;
    BOOLEAN found;

    if (!pmm_ready || count == 0 || count > pmm_total_pages
        || pmm_free_count < count) {
        return 0;
    }

    for (i = 0; i + count <= pmm_total_pages; i++) {
        found = TRUE;
        for (j = 0; j < count; j++) {
            if (pmm_test_bit(i + j)) {
                found = FALSE;
                i += j;            /* skip past the used page */
                break;
            }
        }
        if (found) {
            for (j = 0; j < count; j++) {
                pmm_set_bit(i + j);
            }
            pmm_free_count -= count;
            return pmm_pool_base + ((UINT64)i * PAGE_SIZE);
        }
    }

    return 0;
}

/* Free `count` contiguous pages starting at phys_addr. */
void pmm_free_pages(UINT64 phys_addr, UINTN count)
{
    UINTN i;
    for (i = 0; i < count; i++) {
        pmm_free_page(phys_addr + ((UINT64)i * PAGE_SIZE));
    }
}

UINTN pmm_get_free_pages(void)  { return pmm_free_count;  }
UINTN pmm_get_total_pages(void) { return pmm_total_pages; }

/* ================================================================
 *  Paging  –  manipulate the live 4-level x86-64 page tables
 * ================================================================ */

/* ================================================================
 *  Paging – low-level helpers
 * ================================================================ */

/* Read the CR3 register (physical address of PML4). */
static UINT64 read_cr3(void)
{
    UINT64 cr3;
    __asm__ __volatile__("mov %%cr3, %0" : "=r"(cr3));
    return cr3;
}

/* Invalidate the TLB entry for virtual address `addr`. */
static void invlpg(UINT64 addr)
{
    __asm__ __volatile__("invlpg (%0)" :: "r"(addr) : "memory");
}

/* Return a pointer to the current PML4 table. */
static UINT64 *get_pml4(void)
{
    return (UINT64 *)(UINTN)(read_cr3() & PTE_ADDR_MASK);
}

/*
 * Walk one level of the page table hierarchy.
 * If `create` is TRUE and the entry is missing, a fresh zeroed page is
 * allocated from the PMM and installed.
 */
static UINT64 *paging_walk_level(UINT64 *table, UINTN index, BOOLEAN create)
{
    UINT64 *next;
    UINTN i;
    UINT64 page;

    if (table[index] & PTE_PRESENT) {
        return (UINT64 *)(UINTN)(table[index] & PTE_ADDR_MASK);
    }

    if (!create) {
        return NULL;
    }

    page = pmm_alloc_page();
    if (page == 0) {
        return NULL;
    }

    /* Zero the freshly-allocated page table */
    next = (UINT64 *)(UINTN)page;
    for (i = 0; i < PAGE_SIZE / sizeof(UINT64); i++) {
        next[i] = 0;
    }

    table[index] = page | PTE_PRESENT | PTE_WRITABLE;
    return next;
}

/* ----------------------------------------------------------------
 *  Paging – public interface
 * ---------------------------------------------------------------- */

/* Log the current CR3 value (identity-mapped by UEFI). */
void paging_init(BootInfo *Boot)
{
    SAFE_PRINT(Boot, L"  Page: CR3 = 0x%lx (identity-mapped by loader)\n\r",
               read_cr3());
}

/*
 * Map a single 4 KB page: virt → phys with the given flags.
 * Returns TRUE on success, FALSE if a huge page is in the way or
 * page-table allocation failed.
 */
BOOLEAN paging_map_page(UINT64 virt, UINT64 phys, UINT64 flags)
{
    UINT64 *pml4, *pdpt, *pd, *pt;
    UINTN pml4i, pdpti, pdi, pti;

    pml4i = (virt >> 39) & 0x1FF;
    pdpti = (virt >> 30) & 0x1FF;
    pdi   = (virt >> 21) & 0x1FF;
    pti   = (virt >> 12) & 0x1FF;

    pml4 = get_pml4();

    pdpt = paging_walk_level(pml4, pml4i, TRUE);
    if (pdpt == NULL) return FALSE;

    /* 1 GB huge page – cannot carve a 4 KB mapping inside it */
    if (pdpt[pdpti] & PTE_HUGE) return FALSE;

    pd = paging_walk_level(pdpt, pdpti, TRUE);
    if (pd == NULL) return FALSE;

    /* 2 MB huge page – cannot carve a 4 KB mapping inside it */
    if (pd[pdi] & PTE_HUGE) return FALSE;

    pt = paging_walk_level(pd, pdi, TRUE);
    if (pt == NULL) return FALSE;

    pt[pti] = (phys & PTE_ADDR_MASK) | flags | PTE_PRESENT;
    invlpg(virt);
    return TRUE;
}

/* Remove the mapping for a single 4 KB page and flush the TLB. */
void paging_unmap_page(UINT64 virt)
{
    UINT64 *pml4, *pdpt, *pd, *pt;
    UINTN pml4i, pdpti, pdi, pti;

    pml4i = (virt >> 39) & 0x1FF;
    pdpti = (virt >> 30) & 0x1FF;
    pdi   = (virt >> 21) & 0x1FF;
    pti   = (virt >> 12) & 0x1FF;

    pml4 = get_pml4();

    pdpt = paging_walk_level(pml4, pml4i, FALSE);
    if (pdpt == NULL) return;
    if (pdpt[pdpti] & PTE_HUGE) return;

    pd = paging_walk_level(pdpt, pdpti, FALSE);
    if (pd == NULL) return;
    if (pd[pdi] & PTE_HUGE) return;

    pt = paging_walk_level(pd, pdi, FALSE);
    if (pt == NULL) return;

    pt[pti] = 0;
    invlpg(virt);
}

/*
 * Translate a virtual address to its physical counterpart.
 * Handles 4 KB, 2 MB, and 1 GB page sizes.  Returns 0 if unmapped.
 */
UINT64 paging_get_phys(UINT64 virt)
{
    UINT64 *pml4, *pdpt, *pd, *pt;
    UINTN pml4i, pdpti, pdi, pti;

    pml4i = (virt >> 39) & 0x1FF;
    pdpti = (virt >> 30) & 0x1FF;
    pdi   = (virt >> 21) & 0x1FF;
    pti   = (virt >> 12) & 0x1FF;

    pml4 = get_pml4();

    if (!(pml4[pml4i] & PTE_PRESENT)) return 0;
    pdpt = (UINT64 *)(UINTN)(pml4[pml4i] & PTE_ADDR_MASK);

    if (!(pdpt[pdpti] & PTE_PRESENT)) return 0;
    if (pdpt[pdpti] & PTE_HUGE) {
        /* 1 GB page */
        return (pdpt[pdpti] & 0x000FFFFFC0000000ULL) | (virt & 0x3FFFFFFFULL);
    }
    pd = (UINT64 *)(UINTN)(pdpt[pdpti] & PTE_ADDR_MASK);

    if (!(pd[pdi] & PTE_PRESENT)) return 0;
    if (pd[pdi] & PTE_HUGE) {
        /* 2 MB page */
        return (pd[pdi] & 0x000FFFFFFFE00000ULL) | (virt & 0x1FFFFFULL);
    }
    pt = (UINT64 *)(UINTN)(pd[pdi] & PTE_ADDR_MASK);

    if (!(pt[pti] & PTE_PRESENT)) return 0;
    return (pt[pti] & PTE_ADDR_MASK) | (virt & 0xFFFULL);
}

/* ================================================================
 *  Heap Allocator  –  first-fit free-list with coalescing
 * ================================================================ */

static HeapBlock *heap_start = NULL;
static BOOLEAN    heap_ready = FALSE;

/* Round `val` up to the next multiple of `align`. */
static UINTN align_up(UINTN val, UINTN align)
{
    return (val + align - 1) & ~(align - 1);
}

/*
 * Initialise the heap: allocate HEAP_INITIAL_PAGES from the PMM
 * and set up a single free block spanning the entire region.
 */
void heap_init(BootInfo *Boot)
{
    UINT64 phys;
    UINTN  heap_size;

    phys = pmm_alloc_pages(HEAP_INITIAL_PAGES);
    if (phys == 0) {
        SAFE_PRINT(Boot, L"  Heap: failed to allocate pages\n\r");
        return;
    }

    heap_size  = HEAP_INITIAL_PAGES * PAGE_SIZE;
    heap_start = (HeapBlock *)(UINTN)phys;

    heap_start->magic = HEAP_BLOCK_MAGIC;
    heap_start->state = HEAP_BLOCK_FREE;
    heap_start->size  = heap_size - sizeof(HeapBlock);
    heap_start->next  = NULL;
    heap_start->prev  = NULL;

    heap_ready = TRUE;

    SAFE_PRINT(Boot, L"  Heap: %d KB at 0x%lx\n\r",
               heap_size / 1024, phys);
}

/*
 * Allocate `size` bytes from the heap (first-fit).
 * The returned pointer is aligned to HEAP_ALIGN.  Returns NULL on
 * failure or heap corruption.
 */
void *kmalloc(UINTN size)
{
    HeapBlock *block, *split;
    UINTN aligned;

    if (!heap_ready || size == 0) {
        return NULL;
    }

    aligned = align_up(size, HEAP_ALIGN);

    for (block = heap_start; block != NULL; block = block->next) {
        if (block->magic != HEAP_BLOCK_MAGIC) {
            return NULL;            /* heap corruption */
        }

        if (block->state != HEAP_BLOCK_FREE || block->size < aligned) {
            continue;
        }

        /* Try to split if there is room for another header + 16 bytes */
        if (block->size >= aligned + sizeof(HeapBlock) + HEAP_ALIGN) {
            split = (HeapBlock *)((UINT8 *)block + sizeof(HeapBlock) + aligned);
            split->magic = HEAP_BLOCK_MAGIC;
            split->state = HEAP_BLOCK_FREE;
            split->size  = block->size - aligned - sizeof(HeapBlock);
            split->next  = block->next;
            split->prev  = block;

            if (block->next != NULL) {
                block->next->prev = split;
            }

            block->next = split;
            block->size = aligned;
        }

        block->state = HEAP_BLOCK_USED;
        return (void *)((UINT8 *)block + sizeof(HeapBlock));
    }

    return NULL;                    /* out of heap memory */
}

/*
 * Free a previously kmalloc'd pointer.  Coalesces adjacent free
 * blocks to reduce fragmentation.
 */
void kfree(void *ptr)
{
    HeapBlock *block;

    if (ptr == NULL || !heap_ready) {
        return;
    }

    block = (HeapBlock *)((UINT8 *)ptr - sizeof(HeapBlock));

    if (block->magic != HEAP_BLOCK_MAGIC || block->state != HEAP_BLOCK_USED) {
        return;                     /* bad pointer or double-free */
    }

    block->state = HEAP_BLOCK_FREE;

    /* Coalesce with next neighbour */
    if (block->next != NULL
        && block->next->magic == HEAP_BLOCK_MAGIC
        && block->next->state == HEAP_BLOCK_FREE) {
        block->size += sizeof(HeapBlock) + block->next->size;
        block->next  = block->next->next;
        if (block->next != NULL) {
            block->next->prev = block;
        }
    }

    /* Coalesce with previous neighbour */
    if (block->prev != NULL
        && block->prev->magic == HEAP_BLOCK_MAGIC
        && block->prev->state == HEAP_BLOCK_FREE) {
        block->prev->size += sizeof(HeapBlock) + block->size;
        block->prev->next  = block->next;
        if (block->next != NULL) {
            block->next->prev = block->prev;
        }
    }
}

/* Gather aggregate heap statistics. */
void heap_get_stats(UINTN *total, UINTN *used, UINTN *free_mem,
                    UINTN *num_blocks)
{
    HeapBlock *b;

    *total = 0;  *used = 0;  *free_mem = 0;  *num_blocks = 0;

    if (!heap_ready) return;

    for (b = heap_start; b != NULL && b->magic == HEAP_BLOCK_MAGIC;
         b = b->next) {
        (*num_blocks)++;
        *total += b->size;
        if (b->state == HEAP_BLOCK_USED) {
            *used += b->size;
        } else {
            *free_mem += b->size;
        }
    }
}

/* ================================================================
 *  Top-level helpers
 * ================================================================ */

/* Initialise all memory subsystems in order. */
void memory_init(BootInfo *Boot)
{
    SAFE_PRINT(Boot, L"Initializing memory management...\n\r");
    pmm_init(Boot);
    paging_init(Boot);
    heap_init(Boot);
    SAFE_PRINT(Boot, L"Memory management ready.\n\r\n\r");
}

/* Print a summary of PMM, heap, and paging state to the console. */
void memory_print_stats(BootInfo *Boot)
{
    UINTN h_total, h_used, h_free, h_blocks;
    UINTN p_total, p_free, p_used;
    Task *caller;

    /* Subsystem-level privilege enforcement: memory stats require KERNEL. */
    caller = task_current();
    if (caller != NULL && task_get_privilege(caller) < TASK_PRIV_KERNEL) {
        SAFE_PRINT(Boot, L"Permission denied: memory stats require kernel privilege.\n\r");
        return;
    }

    p_total = pmm_get_total_pages();
    p_free  = pmm_get_free_pages();
    p_used  = p_total - p_free;

    heap_get_stats(&h_total, &h_used, &h_free, &h_blocks);

    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Memory Statistics\n\r");
    SAFE_PRINT(Boot, L"================================================\n\r");
    SAFE_PRINT(Boot, L"\n\r");

    SAFE_PRINT(Boot, L"Physical Memory Manager:\n\r");
    SAFE_PRINT(Boot, L"  Pool Base:   0x%lx\n\r", pmm_pool_base);
    SAFE_PRINT(Boot, L"  Total Pages: %d (%d KB)\n\r",
               p_total, (p_total * PAGE_SIZE) / 1024);
    SAFE_PRINT(Boot, L"  Used Pages:  %d (%d KB)\n\r",
               p_used, (p_used * PAGE_SIZE) / 1024);
    SAFE_PRINT(Boot, L"  Free Pages:  %d (%d KB)\n\r",
               p_free, (p_free * PAGE_SIZE) / 1024);
    SAFE_PRINT(Boot, L"\n\r");

    SAFE_PRINT(Boot, L"Heap Allocator:\n\r");
    SAFE_PRINT(Boot, L"  Total:  %d bytes\n\r", h_total);
    SAFE_PRINT(Boot, L"  Used:   %d bytes\n\r", h_used);
    SAFE_PRINT(Boot, L"  Free:   %d bytes\n\r", h_free);
    SAFE_PRINT(Boot, L"  Blocks: %d\n\r", h_blocks);
    SAFE_PRINT(Boot, L"\n\r");

    SAFE_PRINT(Boot, L"Paging:\n\r");
    SAFE_PRINT(Boot, L"  CR3:  0x%lx\n\r", read_cr3());
    SAFE_PRINT(Boot, L"  Mode: 4-level (PML4)\n\r");
    SAFE_PRINT(Boot, L"\n\r");
}