/*
 * ELSA OS - Virtual Memory Manager
 * Temporal Epistemic Separation Kernel v0.4
 * Phase 1: Paging + Ring3 + APIC + TSC/HPET
 * 2026-05-16
 */

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include "vmm.h"
#include "../arch/x86_64/paging.h"

/* ============================================================
 * FORWARD DECLARATIONS (internal helpers)
 * ============================================================ */

static kernel_pt_t *vmm_alloc_page_table(void);
static void vmm_free_page_table_recursive(void *table, int level);
static uintptr_t vmm_translate_virtual(address_space_t *as, uintptr_t virt);
static pid_t vmm_allocate_pid(void);
static uintptr_t vmm_alloc_physical_page(void);
static void vmm_free_physical_page(uintptr_t phys);
static void vmm_oom_kill(void);
static void vmm_fatal_fault(uintptr_t addr, uint64_t error_code);

/* ============================================================
 * GLOBAL STATE
 * ============================================================ */

static kernel_page_table_t *kernel_pml4 = NULL;
static bitmap_t *kernel_vmm_bitmap = NULL;
static uintptr_t kernel_heap_start = 0;
static uintptr_t kernel_heap_end = 0;
static size_t kernel_heap_used = 0;

/* ============================================================
 * PML4 INITIALIZATION
 * ============================================================ */

int vmm_init(uintptr_t phys_mem_start, size_t phys_mem_size,
             uintptr_t kernel_virt_base)
{
    /* Allocate PML4 (4KB aligned) */
    kernel_pml4 = (kernel_page_table_t *)phys_mem_start;
    memset(kernel_pml4, 0, sizeof(kernel_page_table_t));

    /* Identity map kernel region */
    uintptr_t kernel_phys_start = phys_mem_start + 0x1000; /* after PML4 */
    size_t kernel_size = 0x200000; /* 2MB initial kernel mapping */

    for (size_t i = 0; i < kernel_size; i += PAGE_SIZE_2MB) {
        vmm_map_page(kernel_pml4,
                     kernel_virt_base + i,
                     kernel_phys_start + i,
                     PAGE_FLAG_PRESENT | PAGE_FLAG_WRITE | PAGE_FLAG_HUGE);
    }

    /* Map higher-half kernel space */
    kernel_heap_start = HIGHER_HALF_BASE;
    kernel_heap_end = HIGHER_HALF_BASE + (phys_mem_size / 2);
    kernel_heap_used = 0;

    /* Initialize VMM bitmap for virtual address tracking */
    size_t bitmap_size = (phys_mem_size / PAGE_SIZE) / 8;
    kernel_vmm_bitmap = (bitmap_t *)(kernel_phys_start + kernel_size);
    memset(kernel_vmm_bitmap, 0, bitmap_size);

    /* Load CR3 with PML4 physical address */
    paging_load_cr3((uintptr_t)kernel_pml4);

    /* Enable paging: set PG and WP bits in CR0 */
    paging_enable_paging();

    return 0;
}

/* ============================================================
 * PAGE MAPPING
 * ============================================================ */

int vmm_map_page(kernel_page_table_t *pml4, uintptr_t virt,
                 uintptr_t phys, uint64_t flags)
{
    if (!pml4) return -1;

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

    /* PML4 entry */
    kernel_pdp_t *pdp = NULL;
    if (!(pml4->entries[pml4_idx] & PAGE_FLAG_PRESENT)) {
        pdp = vmm_alloc_page_table();
        if (!pdp) return -1;
        memset(pdp, 0, sizeof(kernel_pdp_t));
        pml4->entries[pml4_idx] = (uintptr_t)pdp | PAGE_FLAG_PRESENT
                                                | PAGE_FLAG_WRITE
                                                | PAGE_FLAG_USER;
    } else {
        pdp = (kernel_pdp_t *)(pml4->entries[pml4_idx] & PAGE_ADDR_MASK);
    }

    /* PDP entry */
    kernel_pd_t *pd = NULL;
    if (!(pdp->entries[pdp_idx] & PAGE_FLAG_PRESENT)) {
        pd = vmm_alloc_page_table();
        if (!pd) return -1;
        memset(pd, 0, sizeof(kernel_pd_t));
        pdp->entries[pdp_idx] = (uintptr_t)pd | PAGE_FLAG_PRESENT
                                              | PAGE_FLAG_WRITE
                                              | PAGE_FLAG_USER;
    } else {
        pd = (kernel_pd_t *)(pdp->entries[pdp_idx] & PAGE_ADDR_MASK);
    }

    /* PD entry - check for 2MB huge page */
    if (flags & PAGE_FLAG_HUGE) {
        pd->entries[pd_idx] = (phys & PAGE_ADDR_MASK_2MB) | flags;
        paging_invalidate_tlb(virt);
        return 0;
    }

    /* PT entry */
    kernel_pt_t *pt = NULL;
    if (!(pd->entries[pd_idx] & PAGE_FLAG_PRESENT)) {
        pt = vmm_alloc_page_table();
        if (!pt) return -1;
        memset(pt, 0, sizeof(kernel_pt_t));
        pd->entries[pd_idx] = (uintptr_t)pt | PAGE_FLAG_PRESENT
                                            | PAGE_FLAG_WRITE
                                            | PAGE_FLAG_USER;
    } else {
        pt = (kernel_pt_t *)(pd->entries[pd_idx] & PAGE_ADDR_MASK);
    }

    pt->entries[pt_idx] = (phys & PAGE_ADDR_MASK) | flags;
    paging_invalidate_tlb(virt);
    return 0;
}

int vmm_unmap_page(kernel_page_table_t *pml4, uintptr_t virt)
{
    if (!pml4) return -1;

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

    kernel_pdp_t *pdp = (kernel_pdp_t *)(pml4->entries[pml4_idx] & PAGE_ADDR_MASK);
    if (!pdp || !(pml4->entries[pml4_idx] & PAGE_FLAG_PRESENT)) return -1;

    kernel_pd_t *pd = (kernel_pd_t *)(pdp->entries[pdp_idx] & PAGE_ADDR_MASK);
    if (!pd || !(pdp->entries[pdp_idx] & PAGE_FLAG_PRESENT)) return -1;

    kernel_pt_t *pt = (kernel_pt_t *)(pd->entries[pd_idx] & PAGE_ADDR_MASK);
    if (!pt || !(pd->entries[pd_idx] & PAGE_FLAG_PRESENT)) return -1;

    pt->entries[pt_idx] = 0;
    paging_invalidate_tlb(virt);
    return 0;
}

/* ============================================================
 * ADDRESS SPACE MANAGEMENT
 * ============================================================ */

address_space_t *vmm_create_address_space(void)
{
    address_space_t *as = (address_space_t *)vmm_kernel_alloc(sizeof(address_space_t));
    if (!as) return NULL;

    as->pml4 = (kernel_page_table_t *)vmm_alloc_page_table();
    if (!as->pml4) {
        vmm_kernel_free(as, sizeof(address_space_t));
        return NULL;
    }
    memset(as->pml4, 0, sizeof(kernel_page_table_t));

    as->base_virt = USER_SPACE_BASE;
    as->size = USER_SPACE_SIZE;
    as->pid = vmm_allocate_pid();
    as->refcount = 1;

    return as;
}

int vmm_switch_address_space(address_space_t *as)
{
    if (!as || !as->pml4) return -1;
    paging_load_cr3((uintptr_t)as->pml4);
    return 0;
}

void vmm_destroy_address_space(address_space_t *as)
{
    if (!as) return;
    as->refcount--;
    if (as->refcount > 0) return;

    /* Free all page tables recursively */
    vmm_free_page_table_recursive(as->pml4, 4);
    vmm_kernel_free(as, sizeof(address_space_t));
}

/* ============================================================
 * COPY-ON-WRITE
 * ============================================================ */

int vmm_copy_on_write(address_space_t *src, address_space_t *dst,
                      uintptr_t virt, size_t size)
{
    for (uintptr_t v = virt; v < virt + size; v += PAGE_SIZE) {
        uint64_t phys = vmm_translate_virtual(src, v);
        if (phys == 0) continue;

        /* Map as read-only in both spaces */
        vmm_map_page(dst->pml4, v, phys,
                     PAGE_FLAG_PRESENT | PAGE_FLAG_USER);
        vmm_map_page(src->pml4, v, phys,
                     PAGE_FLAG_PRESENT | PAGE_FLAG_USER);
    }
    return 0;
}

/* ============================================================
 * PAGE FAULT HANDLER
 * ============================================================ */

void vmm_page_fault_handler(uintptr_t fault_addr, uint64_t error_code)
{
    bool present    = (error_code & 0x1) != 0;
    bool write      = (error_code & 0x2) != 0;
    bool user       = (error_code & 0x4) != 0;
    bool reserved   = (error_code & 0x8) != 0;
    bool exec       = (error_code & 0x10) != 0;

    /* Copy-on-write fault */
    if (present && write) {
        /* Allocate new page, copy data, update mapping */
        uintptr_t new_phys = vmm_alloc_physical_page();
        if (new_phys == 0) {
            vmm_oom_kill();
            return;
        }
        /* TODO: copy old page content to new page */
        /* TODO: update PTE to point to new_phys with write flag */
        return;
    }

    /* Demand paging */
    if (!present && !reserved) {
        /* TODO: load page from swap/disk */
        return;
    }

    /* Fatal fault */
    vmm_fatal_fault(fault_addr, error_code);
}

/* ============================================================
 * KERNEL ALLOCATOR
 * ============================================================ */

void *vmm_kernel_alloc(size_t size)
{
    size = (size + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);

    if (kernel_heap_used + size > (kernel_heap_end - kernel_heap_start)) {
        return NULL; /* OOM */
    }

    void *ptr = (void *)(kernel_heap_start + kernel_heap_used);
    kernel_heap_used += size;

    /* Zero the allocation */
    memset(ptr, 0, size);
    return ptr;
}

void vmm_kernel_free(void *ptr, size_t size)
{
    /* Simple bump allocator - no individual free */
    (void)ptr;
    (void)size;
}

/* ============================================================
 * INTERNAL HELPERS
 * ============================================================ */

static kernel_pt_t *vmm_alloc_page_table(void)
{
    uintptr_t phys = vmm_alloc_physical_page();
    if (phys == 0) return NULL;
    return (kernel_pt_t *)phys;
}

static void vmm_free_page_table_recursive(void *table, int level)
{
    if (!table || level <= 0) return;

    uint64_t *entries = (uint64_t *)table;
    for (int i = 0; i < 512; i++) {
        if (entries[i] & PAGE_FLAG_PRESENT) {
            void *child = (void *)(entries[i] & PAGE_ADDR_MASK);
            if (level > 1) {
                vmm_free_page_table_recursive(child, level - 1);
            }
        }
    }

    vmm_free_physical_page((uintptr_t)table);
}

static uintptr_t vmm_translate_virtual(address_space_t *as, uintptr_t virt)
{
    if (!as || !as->pml4) return 0;

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

    kernel_pdp_t *pdp = (kernel_pdp_t *)(as->pml4->entries[pml4_idx] & PAGE_ADDR_MASK);
    if (!pdp) return 0;

    kernel_pd_t *pd = (kernel_pd_t *)(pdp->entries[pdp_idx] & PAGE_ADDR_MASK);
    if (!pd) return 0;

    /* Check for huge page */
    if (pd->entries[pd_idx] & PAGE_FLAG_HUGE) {
        return (pd->entries[pd_idx] & PAGE_ADDR_MASK_2MB) | (virt & 0x1FFFFF);
    }

    kernel_pt_t *pt = (kernel_pt_t *)(pd->entries[pd_idx] & PAGE_ADDR_MASK);
    if (!pt) return 0;

    return (pt->entries[pt_idx] & PAGE_ADDR_MASK) | (virt & 0xFFF);
}

static pid_t vmm_allocate_pid(void)
{
    static pid_t next_pid = 1;
    return next_pid++;
}

static uintptr_t vmm_alloc_physical_page(void)
{
    /* TODO: integrate with physical memory manager */
    static uintptr_t next_phys = 0x100000; /* 1MB */
    uintptr_t page = next_phys;
    next_phys += PAGE_SIZE;
    return page;
}

static void vmm_free_physical_page(uintptr_t phys)
{
    /* TODO: return page to physical memory manager */
    (void)phys;
}

static void vmm_oom_kill(void)
{
    /* ELSA OS: ABSTAIN - no execution under OOM */
    /* TODO: invoke temporal decision kernel for OOM handling */
    while (1) { __asm__ volatile("hlt"); }
}

static void vmm_fatal_fault(uintptr_t addr, uint64_t error_code)
{
    /* ELSA OS: ABSTAIN - no execution under fatal fault */
    /* TODO: log to AuditChain, halt core */
    (void)addr;
    (void)error_code;
    while (1) { __asm__ volatile("hlt"); }
}