#ifndef MEMORY_H
#define MEMORY_H

#include <efi.h>
#include "boot_info.h"

/* Page size constants */
#define PAGE_SIZE       4096
#define PAGE_SHIFT      12
#define PAGE_MASK       (~(UINT64)(PAGE_SIZE - 1))

/* Page table entry flags */
#define PTE_PRESENT     (1ULL << 0)
#define PTE_WRITABLE    (1ULL << 1)
#define PTE_USER        (1ULL << 2)
#define PTE_WRITETHROUGH (1ULL << 3)
#define PTE_NOCACHE     (1ULL << 4)
#define PTE_ACCESSED    (1ULL << 5)
#define PTE_DIRTY       (1ULL << 6)
#define PTE_HUGE        (1ULL << 7)
#define PTE_GLOBAL      (1ULL << 8)
#define PTE_NX          (1ULL << 63)

#define PTE_ADDR_MASK   0x000FFFFFFFFFF000ULL

/* PMM pool: 16 MB managed by the kernel */
#define PMM_POOL_SIZE   (16ULL * 1024 * 1024)
#define PMM_POOL_PAGES  (PMM_POOL_SIZE / PAGE_SIZE)

/* Heap constants */
#define HEAP_INITIAL_PAGES  64   /* 256 KB initial heap */
#define HEAP_BLOCK_MAGIC    0xDEADBEEFU
#define HEAP_BLOCK_FREE     0
#define HEAP_BLOCK_USED     1
#define HEAP_ALIGN          16

/* -------- Physical Memory Manager -------- */

void    pmm_init(BootInfo *Boot);
UINT64  pmm_alloc_page(void);
void    pmm_free_page(UINT64 phys_addr);
UINT64  pmm_alloc_pages(UINTN count);
void    pmm_free_pages(UINT64 phys_addr, UINTN count);
UINTN   pmm_get_free_pages(void);
UINTN   pmm_get_total_pages(void);

/* -------- Paging -------- */

void    paging_init(BootInfo *Boot);
BOOLEAN paging_map_page(UINT64 virt, UINT64 phys, UINT64 flags);
void    paging_unmap_page(UINT64 virt);
UINT64  paging_get_phys(UINT64 virt);

/* -------- Heap Allocator -------- */

typedef struct HeapBlock {
    UINT32 magic;
    UINT32 state;
    UINTN  size;              /* usable bytes (excludes header) */
    struct HeapBlock *next;
    struct HeapBlock *prev;
} HeapBlock;

void  heap_init(BootInfo *Boot);
void *kmalloc(UINTN size);
void  kfree(void *ptr);
void  heap_get_stats(UINTN *total, UINTN *used, UINTN *free_mem,
                     UINTN *num_blocks);

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

void memory_init(BootInfo *Boot);
void memory_print_stats(BootInfo *Boot);

#endif