/*
 * Universalisos Memory Management Implementation - Stage 3
 * PikeOS 5.0 Feature Parity - Memory Virtualization
 *
 * Stage 3: Memory Management + VM Context Switching
 * Author: PortugalFuturista Hypervisor Development Team
 * Version: 1.0.0
 */

#include "mm.h"
#include "arch/arm/uart.h"
#include <stdint.h>

// Global memory management state
static mm_state_t g_mm_state = {
    .kernel_page_table = {
        .ttb0 = nullptr,
        .ttb1 = nullptr,
        .ttb0_size = 0,
        .ttb1_size = 0,
        .using_ttbr0 = false,
        .current_domain = 0
    },
    .vm_page_tables = {},
    .active_vm_count = 0,
    .domains = {},
    .active_domain_count = 0,
    .total_memory = 0,
    .free_memory = 0,
    .hypervisor_memory = 0,
    .mmu_enabled = false,
    .caches_enabled = false
};

// Simple page table allocator (Stage 3 - will be improved later)
static uint32_t page_table_memory[4096] __attribute__((aligned(16384))); // 16KB aligned

/* Expose the flat kernel pgdir so D-1's per-VM cloned pgdirs can copy it. */
extern "C" uint32_t *uos_get_kernel_pgdir(void) { return page_table_memory; }

/**
 * Initialize Memory Management System
 */
extern "C" void mm_init(void) {
    uart_puts("MM: Initializing Memory Management System\n");

    // Initialize state
    g_mm_state.total_memory = 512 * 1024 * 1024; // 512MB for QEMU virt
    g_mm_state.free_memory = g_mm_state.total_memory - (2 * 1024 * 1024); // Reserve 2MB for hypervisor
    g_mm_state.hypervisor_memory = 2 * 1024 * 1024;
    g_mm_state.mmu_enabled = false;
    g_mm_state.caches_enabled = false;

    // Initialize VM page tables
    for (int i = 0; i < 16; i++) {
        g_mm_state.vm_page_tables[i].ttb0 = nullptr;
        g_mm_state.vm_page_tables[i].ttb1 = nullptr;
        g_mm_state.vm_page_tables[i].ttb0_size = 0;
        g_mm_state.vm_page_tables[i].ttb1_size = 0;
    }

    g_mm_state.active_vm_count = 0;

    uart_puts("MM: Memory Management System initialized\n");
    uart_puts("MM: Total memory: ");
    uart_print_dec(g_mm_state.total_memory / (1024 * 1024));
    uart_puts(" MB\n");

    uart_puts("MM: Free memory: ");
    uart_print_dec(g_mm_state.free_memory / (1024 * 1024));
    uart_puts(" MB\n");
}

/**
 * Initialize ARMv7 MMU
 */
extern "C" void mmu_init(void) {
    uart_puts("MMU: Initializing ARMv7 MMU\n");

    // Initialize kernel page table with simple identity mapping
    g_mm_state.kernel_page_table.ttb0 = page_table_memory;
    g_mm_state.kernel_page_table.ttb0_size = 16384; // 16KB
    g_mm_state.kernel_page_table.ttb1 = nullptr;
    g_mm_state.kernel_page_table.using_ttbr0 = true;
    g_mm_state.kernel_page_table.current_domain = 0; // Domain 0

    // Clear page table
    for (int i = 0; i < 4096; i++) {
        page_table_memory[i] = 0;
    }

    // PikeOS-style flat 1:1 section map (non-LPAE ARMv7), covering the full 4 GB
    // address space so every physical region is reachable after the MMU is on.
    // Adapted from PikeOS boot_map.c (psp_arm_boot_map_ram / _io). Constants use
    // the Universalisos uos_ convention instead of PikeOS's p4_/PD_.
    //
    // Section descriptor layout (armmmu-v6.h):
    //   [31:20] base   [19:12] TEX   [11:10] AP   [9] impl   [8:5] domain/impl/AP2
    //   [4] XN   [3] C   [2] B   [1] section-type(=1)   [0] 0
    //   => section type = 0b10 => bit1 set.
    #define UOS_PD_SECT   (1u << 1)   /* section descriptor */
    #define UOS_PD_B      (1u << 2)   /* bufferable */
    #define UOS_PD_C      (1u << 3)   /* cacheable */
    #define UOS_PD_XN     (1u << 4)   /* execute never */
    #define UOS_PD_AP0    (1u << 10)  /* AP[1:0] = 01 => kernel RW, user none */
    #define UOS_PD_TEX0   (1u << 12)  /* TEX bit 0 */

    /* RAM: cacheable normal write-back (TEX=001,C=1,B=1), executable.
     * I/O / device: Device memory (TEX=000,C=0,B=1), executable.
     *
     * Device memory (not strongly-ordered) is the correct attribute for MMIO
     * registers including PCIe ECAM: it permits write buffering while keeping
     * device access semantics, and avoids the strict completion-ack demands of
     * strongly-ordered memory under which QEMU's gpex ECAM reads can deadlock.
     * The whole map stays executable so the exception vectors at 0x0 are
     * reachable (matches PikeOS's boot identity map). */
    const uint32_t ram_section = UOS_PD_SECT | UOS_PD_AP0 | UOS_PD_TEX0 | UOS_PD_C | UOS_PD_B;
    const uint32_t io_section  = UOS_PD_SECT | UOS_PD_AP0 | UOS_PD_B;

    /* QEMU virt RAM window: 0x40000000 .. 0x60000000 (512 MiB, sections 1024..1535) */
    const uint32_t ram_first_section = 0x40000000u >> 20;   /* 1024 */
    const uint32_t ram_last_section  = 0x60000000u >> 20;   /* 1536 (exclusive) */

    for (uint32_t i = 0; i < 4096u; i++) {
        uint32_t phys = i << 20;  /* 1 MB section base */
        if (i >= ram_first_section && i < ram_last_section) {
            page_table_memory[i] = phys | ram_section;
        } else {
            page_table_memory[i] = phys | io_section;
        }
    }

    uart_puts("MMU: flat 4GB section map built (RAM cacheable, I/O strongly-ordered)\n");

    /* Turn the MMU on with the flat map we just built (PikeOS-style). With the
     * MMU actually enabled, device regions are properly attributed (strongly
     * ordered) and reachable, and the hypervisor can later add per-VM mappings.
     * This is the fix for the MMIO-access hangs: previously mmu_enable() was
     * never called, leaving the page table built but inert. */
    extern void mmu_enable(void);
    mmu_enable();

    /* D-1: allocate the per-VM page-directory pool (cloned from this flat map). */
    extern int uos_arm_init_mmu(void);
    uos_arm_init_mmu();
}

/**
 * Enable ARMv7 MMU (PikeOS-style, adapted).
 *
 * Mirrors PREBOOT_psp_arm_boot_map_activate() + the startup M-bit flip in
 * PikeOS's boot_map.c: program DACR/TTBCR/TTBR0, flush TLB, then set SCTLR.M.
 * The page table is a flat identity map (VA == PA), so the instruction stream
 * keeps running unchanged across the MMU-on transition.
 *
 * Caches are left OFF on first enable so the page walker always reads the table
 * we just wrote straight from RAM (no D-cache coherency window). This is the
 * conservative, PikeOS-boot-equivalent path.
 */
extern "C" void mmu_enable(void) {
    uart_puts("MMU: enabling (PikeOS-style flat map)\n");

    /* TTB_FLAGS = outer-cacheable write-allocate | inner-region bit0
     * (armmmu-v6.h: TTB_OC_WA | TTB_IRGN0). Harmless with D-cache off. */
    const uint32_t UOS_TTB_FLAGS = (1u << 3) | (1u << 6);

    /* Domain access control: 0x55555555 => every domain = client (01),
     * so the AP bits in each descriptor are checked. */
    __asm__ volatile("mcr p15, 0, %0, c3, c0, 0" : : "r"(0x55555555u));

    /* TTBCR = 0: no TTBR0/TTBR1 split, TTBR0 covers the whole space. */
    __asm__ volatile("mcr p15, 0, %0, c2, c0, 2" : : "r"(0u));

    /* TTBR0 = page table base | TTB_FLAGS. */
    uint32_t ttbr0 = (uint32_t)page_table_memory | UOS_TTB_FLAGS;
    __asm__ volatile("mcr p15, 0, %0, c2, c0, 0" : : "r"(ttbr0));

    /* Push the table writes and drop any stale TLB entries before flipping M. */
    __asm__ volatile("dsb");
    __asm__ volatile("mcr p15, 0, %0, c8, c7, 0" : : "r"(0)); /* TLBIALL */
    __asm__ volatile("isb");

    /* Flip the M bit. Keep A=0 (EABI assumption). Leave C/I as-is. */
    uint32_t sctlr;
    __asm__ volatile("mrc p15, 0, %0, c1, c0, 0" : "=r"(sctlr));
    sctlr |= (1u << 0);    /* M: MMU enable */
    sctlr &= ~(1u << 1);   /* A: alignment checking off (EABI codegen) */
    __asm__ volatile("mcr p15, 0, %0, c1, c0, 0" : : "r"(sctlr));

    __asm__ volatile("dsb");
    __asm__ volatile("isb");

    g_mm_state.mmu_enabled = true;

    uart_puts("MMU: enabled, identity map active\n");
}

/**
 * Disable ARMv7 MMU
 */
extern "C" void mmu_disable(void) {
    uart_puts("MMU: Disabling MMU\n");

    uint32_t sctlr;
    __asm__ volatile("mrc p15, 0, %0, c1, c0, 0" : "=r"(sctlr));

    sctlr &= ~(1 << 0);  // M bit - MMU disable
    sctlr &= ~(1 << 2);  // C bit - Data cache disable
    sctlr &= ~(1 << 12); // I bit - Instruction cache disable

    __asm__ volatile("mcr p15, 0, %0, c1, c0, 0" : : "r"(sctlr));

    g_mm_state.mmu_enabled = false;
    g_mm_state.caches_enabled = false;

    uart_puts("MMU: MMU and caches disabled\n");
}

/**
 * Create page table for VM
 */
extern "C" page_table_t* mm_create_page_table(uint32_t vm_id) {
    if (vm_id >= 16) {
        uart_puts("MM: Invalid VM ID for page table creation\n");
        return nullptr;
    }

    uart_puts("MM: Creating page table for VM ");
    uart_print_dec(vm_id);
    uart_puts("\n");

    // Allocate page table (simplified for Stage 3)
    static uint32_t vm_page_tables[16][4096] __attribute__((aligned(16384)));

    page_table_t* pt = &g_mm_state.vm_page_tables[vm_id];
    pt->ttb0 = vm_page_tables[vm_id];
    pt->ttb0_size = 16384;
    pt->ttb1 = nullptr;
    pt->using_ttbr0 = true;
    pt->current_domain = vm_id + 1; // Use VM ID + 1 as domain

    // Clear page table
    for (int i = 0; i < 4096; i++) {
        vm_page_tables[vm_id][i] = 0;
    }

    g_mm_state.active_vm_count++;

    uart_puts("MM: Page table created for VM ");
    uart_print_dec(vm_id);
    uart_puts(" (Domain ");
    uart_print_dec(pt->current_domain);
    uart_puts(")\n");

    return pt;
}

/**
 * Destroy page table
 */
extern "C" void mm_destroy_page_table(page_table_t* pt) {
    if (!pt) return;

    uart_puts("MM: Destroying page table\n");

    // Clear page table entries
    if (pt->ttb0) {
        for (int i = 0; i < 4096; i++) {
            ((uint32_t*)pt->ttb0)[i] = 0;
        }
    }

    pt->ttb0 = nullptr;
    pt->ttb0_size = 0;
    pt->using_ttbr0 = false;

    g_mm_state.active_vm_count--;

    uart_puts("MM: Page table destroyed\n");
}

/**
 * Map physical to virtual address
 */
extern "C" bool mm_map_page(page_table_t* pt, uint32_t virt_addr, uint32_t phys_addr,
                           uint32_t permissions, uint32_t domain) {
    (void)permissions; // Unused in Stage 3
    if (!pt || !pt->ttb0) {
        return false;
    }

    // Simple section mapping (1MB granularity for Stage 3)
    uint32_t section_index = (virt_addr >> 20) & 0xFFF;
    uint32_t section_phys = phys_addr & 0xFFF00000;

    uint32_t descriptor = section_phys | 0x400 | 0x10 | 0x02; // Section, domain 0, manager
    descriptor |= (domain & 0xF) << 5; // Set domain

    ((uint32_t*)pt->ttb0)[section_index] = descriptor;

    return true;
}

/**
 * Unmap virtual address
 */
extern "C" void mm_unmap_page(page_table_t* pt, uint32_t virt_addr) {
    if (!pt || !pt->ttb0) {
        return;
    }

    uint32_t section_index = (virt_addr >> 20) & 0xFFF;
    ((uint32_t*)pt->ttb0)[section_index] = 0;
}

/**
 * Create memory domain
 */
extern "C" memory_domain_t* mm_create_domain(uint32_t domain_id, const char* name) {
    if (domain_id >= 16) {
        uart_puts("MM: Invalid domain ID\n");
        return nullptr;
    }

    uart_puts("MM: Creating memory domain '");
    uart_puts(name);
    uart_puts("' (ID ");
    uart_print_dec(domain_id);
    uart_puts(")\n");

    memory_domain_t* domain = &g_mm_state.domains[domain_id];
    domain->domain_id = domain_id;
    domain->domain_name = name;
    domain->regions = nullptr;
    domain->region_count = 0;
    domain->access_permissions = DOMAIN_CLIENT;
    domain->enable_protection = true;
    domain->enable_logging = false;

    uart_puts("MM: Memory domain created\n");
    return domain;
}

/**
 * Destroy memory domain
 */
extern "C" void mm_destroy_domain(memory_domain_t* domain) {
    if (!domain) return;

    uart_puts("MM: Destroying memory domain '");
    uart_puts(domain->domain_name);
    uart_puts("'\n");

    domain->domain_id = 0;
    domain->domain_name = nullptr;
    domain->regions = nullptr;
    domain->region_count = 0;

    uart_puts("MM: Memory domain destroyed\n");
}

/**
 * Add memory region to domain
 */
extern "C" bool mm_add_region_to_domain(memory_domain_t* domain, memory_region_t* region) {
    if (!domain || !region) {
        return false;
    }

    uart_puts("MM: Adding region to domain '");
    uart_puts(domain->domain_name);
    uart_puts("'\n");

    // For Stage 3, just print the information
    // In full implementation, would add to domain's region array
    uart_puts("MM:   Base: 0x");
    uart_print_hex(region->base_address);
    uart_puts("\n");

    uart_puts("MM:   Size: ");
    uart_print_dec(region->size / (1024 * 1024));
    uart_puts(" MB\n");

    return true;
}

/**
 * Handle data abort (memory fault)
 */
extern "C" void mm_handle_data_abort(uint32_t fault_address, uint32_t fault_status) {
    uart_puts("\n!!! MM DATA ABORT !!!\n");
    uart_puts("Fault Address: 0x");
    uart_print_hex(fault_address);
    uart_puts("\n");

    uart_puts("Fault Status: 0x");
    uart_print_hex(fault_status);
    uart_puts("\n");

    // Decode fault status
    uart_puts("Fault Type: ");
    if (fault_status & 0x08) uart_puts("Debug event\n");
    else if (fault_status & 0x04) uart_puts("Translation fault\n");
    else if (fault_status & 0x02) uart_puts("Access flag fault\n");
    else if (fault_status & 0x01) uart_puts("Domain fault\n");
    else uart_puts("Unknown\n");

    // For Stage 3, just halt
    uart_puts("System halted for memory safety analysis.\n");
    while(1) {
        __asm__("wfi");
    }
}

/**
 * Handle prefetch abort (instruction fetch fault)
 */
extern "C" void mm_handle_prefetch_abort(uint32_t fault_address, uint32_t fault_status) {
    uart_puts("\n!!! MM PREFETCH ABORT !!!\n");
    uart_puts("Fault Address: 0x");
    uart_print_hex(fault_address);
    uart_puts("\n");

    uart_puts("Fault Status: 0x");
    uart_print_hex(fault_status);
    uart_puts("\n");

    // For Stage 3, just halt
    uart_puts("System halted for memory safety analysis.\n");
    while(1) {
        __asm__("wfi");
    }
}

/**
 * Flush TLB entries
 */
extern "C" void mm_flush_tlb(void) {
    __asm__ volatile("mcr p15, 0, r0, c8, c7, 0"); // TLBIALL
    __asm__ volatile("dsb");
    __asm__ volatile("isb");
}

/**
 * Print memory statistics
 */
extern "C" void mm_print_statistics(void) {
    uart_puts("\n=== Memory Management Statistics ===\n");

    uart_puts("Total Memory: ");
    uart_print_dec(g_mm_state.total_memory / (1024 * 1024));
    uart_puts(" MB\n");

    uart_puts("Free Memory: ");
    uart_print_dec(g_mm_state.free_memory / (1024 * 1024));
    uart_puts(" MB\n");

    uart_puts("Hypervisor Memory: ");
    uart_print_dec(g_mm_state.hypervisor_memory / (1024 * 1024));
    uart_puts(" MB\n");

    uart_puts("Active VMs: ");
    uart_print_dec(g_mm_state.active_vm_count);
    uart_puts("\n");

    uart_puts("MMU Enabled: ");
    uart_puts(g_mm_state.mmu_enabled ? "Yes\n" : "No\n");

    uart_puts("Caches Enabled: ");
    uart_puts(g_mm_state.caches_enabled ? "Yes\n" : "No\n");

    uart_puts("======================================\n\n");
}