//! The kernel's page tables and virtual memory manager.
//!
//! Builds our own 4-level page tables and switches CR3 onto them, replacing the
//! firmware's. Unlike the earlier bootstrap this maps with real permissions:
//! RAM is identity-mapped read-write + no-execute, the kernel's own segments get
//! their ELF permissions (code R+X, rodata R, data R+W+NX), and page 0 is left
//! unmapped as a null guard. It also exposes map/unmap for on-demand mapping,
//! which the kernel heap will build on.
//!
//! Everything is 4 KiB pages — precise and simple; the extra table memory is
//! negligible against available RAM.

const boot_handoff = @import("boot-handoff");
const abi = @import("abi");
const io = @import("io.zig");

const page_size = abi.page_size;

// Page-table entry bits.
const present: u64 = 1 << 0;
const writable: u64 = 1 << 1;
const user: u64 = 1 << 2; // U/S: accessible from ring 3 (must be set at every level)
const pwt: u64 = 1 << 3; // page write-through
const pcd: u64 = 1 << 4; // page cache disable (with PWT: strong-uncacheable under the default PAT)
const device_grant: u64 = 1 << 9; // available bit: this leaf maps device MMIO, not RAM — do not reclaim
const no_execute: u64 = 1 << 63;
const address_mask: u64 = 0x000F_FFFF_FFFF_F000;

// ELF segment flags (p_flags).
const pf_x: u32 = 1;
const pf_w: u32 = 2;

// State kept after init so map()/unmap() can serve later callers (e.g. the heap).
var kernel_pml4: u64 = 0;
var alloc_frame: *const fn () ?u64 = undefined;
var free_frame: *const fn (u64) void = undefined; // for tearing down address spaces

/// Set once the kernel is running on its own tables (past the CR3 load in
/// `init`). Before that, the kernel reaches page-table frames through the
/// *loader's* bootstrap physmap, which only covers the low 4 GiB — so every
/// frame allocated for a table during that window must be below 4 GiB. Both the
/// frame allocator and this code scan from low addresses up, so it holds
/// naturally; the assertion in `allocTable` makes a violation loud rather than
/// a silent fault. After the switch the kernel's own physmap covers all RAM.
var on_own_tables = false;

/// Set at the end of `init`. Guards against a new *higher-half* PML4 entry being
/// created afterward: the kernel half is pre-populated at init and then shared
/// by copying PML4[256..512) into every process address space (M3), so a late
/// top-half entry would be invisible to already-created address spaces.
var init_done = false;

const bootstrap_physmap_limit: u64 = 4 << 30;

/// Dereference a page-table frame by its physical address, via the physmap.
/// This is the single hinge for the higher-half move: page tables hold physical
/// frame addresses (pmm gives out physical frames, and CR3/PTEs must be
/// physical), but the kernel reaches them at `physmap_base + physical`. Valid under
/// both the loader's bootstrap tables and the kernel's own, which share the
/// physmap base.
fn tableAt(physical: u64) *[512]u64 {
    return @ptrFromInt(boot_handoff.physicalToVirtual(physical));
}

fn allocTable() u64 {
    const frame = alloc_frame() orelse @panic("paging: out of memory building page tables");
    if (!on_own_tables and frame >= bootstrap_physmap_limit)
        @panic("paging: table frame above the 4 GiB bootstrap physmap");
    @memset(tableAt(frame)[0..], 0);
    return frame;
}

/// Return the table an entry points at, creating it if empty. Intermediate
/// entries are writable and executable so the leaf's bits govern (a page is
/// writable only if every level is; non-executable if any level is).
fn descend(entry: *u64) u64 {
    if (entry.* & present != 0) return entry.* & address_mask;
    const frame = allocTable();
    entry.* = frame | present | writable;
    return frame;
}

/// Map one 4 KiB page `virtual` -> `physical` with `flags` (present is added).
fn mapPage(pml4: u64, virtual: u64, physical: u64, flags: u64) void {
    const pml4e = &tableAt(pml4)[(virtual >> 39) & 0x1FF];
    // The kernel half is fixed after init: every top-half PML4 entry is
    // pre-created so address spaces can share it by copying these slots. A new
    // one here would be invisible to address spaces already made.
    if (init_done and (virtual >> 63) == 1 and pml4e.* & present == 0)
        @panic("paging: new higher-half PML4 entry after init");
    const pdpt = descend(pml4e);
    const pdpte = &tableAt(pdpt)[(virtual >> 30) & 0x1FF];
    const pd = descend(pdpte);
    const pde = &tableAt(pd)[(virtual >> 21) & 0x1FF];
    const pt = descend(pde);
    tableAt(pt)[(virtual >> 12) & 0x1FF] = (physical & address_mask) | flags | present;
}

/// Map [physical_base, physical_base+len) into the physmap (at physicalToVirtual(physical)) with
/// `flags`, rounded out to whole pages. This is how the kernel keeps a permanent
/// window onto physical memory once the low identity map goes away.
fn mapRangePhysmap(pml4: u64, physical_base: u64, len: u64, flags: u64) void {
    var address = physical_base & ~@as(u64, page_size - 1);
    const end = physical_base + len;
    while (address < end) : (address += page_size) {
        mapPage(pml4, boot_handoff.physicalToVirtual(address), address, flags);
    }
}

fn regions(mm: boot_handoff.MemoryMap) []const boot_handoff.MemoryRegion {
    return @as([*]const boot_handoff.MemoryRegion, @ptrFromInt(boot_handoff.physicalToVirtual(mm.regions)))[0..mm.len];
}

/// Enable the NX bit in the page-table format (EFER.NXE). Must happen before we
/// load a CR3 whose entries set the NX bit, or those bits are reserved and fault.
fn enableNx() void {
    const efer_msr = 0xC0000080;
    io.wrmsr(efer_msr, io.rdmsr(efer_msr) | (1 << 11));
}

/// Build the address space and switch onto it.
pub fn init(allocFrame: *const fn () ?u64, freeFrame: *const fn (u64) void, boot_information: *const boot_handoff.BootInformation) void {
    alloc_frame = allocFrame;
    free_frame = freeFrame;
    enableNx();
    const pml4 = allocTable();

    // 1. All RAM in the physmap (physicalToVirtual(physical)) RW + NX. No identity/low-half
    //    mapping: the low half belongs to user space. MMIO is skipped here and
    //    mapped on demand (mapMmio) or explicitly below.
    for (regions(boot_information.memory_map)) |r| {
        if (r.kind == .mmio) continue;
        mapRangePhysmap(pml4, r.base, r.pages * page_size, present | writable | no_execute);
    }

    // 2. Physmap windows for the framebuffer and the Local APIC (device memory
    //    the kernel touches directly), RW + NX.
    const fb = boot_information.framebuffer;
    mapRangePhysmap(pml4, fb.base, @as(u64, fb.height) * fb.pitch, present | writable | no_execute);
    mapPage(pml4, boot_handoff.physicalToVirtual(0xFEE00000), 0xFEE00000, present | writable | no_execute);

    // 3. The kernel's own segments at their higher-half link addresses, mapped
    //    to their low physical load addresses with real ELF permissions: code
    //    R+X, rodata R, data R+W+NX. This is the W^X guarantee.
    for (boot_information.kernel_segments[0..boot_information.kernel_segment_count]) |seg| {
        var flags: u64 = present;
        if (seg.flags & pf_w != 0) flags |= writable;
        if (seg.flags & pf_x == 0) flags |= no_execute;
        var off: u64 = 0;
        while (off < seg.pages * page_size) : (off += page_size) {
            mapPage(pml4, seg.virtual + off, seg.physical + off, flags);
        }
    }

    // 4. Pre-create every higher-half PML4 entry (an empty PDPT where none
    //    exists yet), so the whole kernel half is a fixed set of top-level
    //    slots. A process address space (M3) then shares the kernel half simply
    //    by copying PML4[256..512) — growth beneath these slots (heap, on-demand
    //    MMIO) propagates to every address space because they share the PDPTs.
    for (256..512) |i| {
        const e = &tableAt(pml4)[i];
        if (e.* & present == 0) e.* = allocTable() | present | writable;
    }

    kernel_pml4 = pml4;
    asm volatile ("mov %[pml4], %%cr3"
        :
        : [pml4] "r" (pml4),
        : .{ .memory = true });
    on_own_tables = true; // now on the kernel's physmap (covers all RAM)
    init_done = true; // the kernel half is fixed from here
}

/// The kernel's own top-level page table (physical). Every kernel task and every
/// per-process address space shares this table's higher half.
pub fn kernelPml4() u64 {
    return kernel_pml4;
}

/// Load CR3 (switch the active address space). `pml4` is a physical frame.
pub fn loadCr3(pml4: u64) void {
    asm volatile ("mov %[pml4], %%cr3"
        :
        : [pml4] "r" (pml4),
        : .{ .memory = true });
}

/// Map a page into the kernel address space on demand (for the heap, etc.).
/// `writable_page` controls W; pages are always mapped non-executable.
pub fn map(virtual: u64, physical: u64, writable_page: bool) void {
    var flags: u64 = present | no_execute;
    if (writable_page) flags |= writable;
    mapPage(kernel_pml4, virtual, physical, flags);
    invalidate(virtual);
}

/// Map a device MMIO range into the physmap and return the virtual address to
/// use for it (physicalToVirtual(physical)). The single way the kernel (and the device
/// layer, via the HAL) reaches memory-mapped registers once the identity map is
/// gone: physmap pages are RW + NX, so a driver never executes device memory.
/// Idempotent for already-mapped ranges. `len` 0 maps one page.
pub fn mapMmio(physical: u64, len: u64, writable_page: bool) u64 {
    var flags: u64 = present | no_execute;
    if (writable_page) flags |= writable;
    const first = physical & ~@as(u64, page_size - 1);
    const last = physical + (if (len == 0) 1 else len) - 1;
    var address = first;
    while (address <= (last & ~@as(u64, page_size - 1))) : (address += page_size) {
        const virtual = boot_handoff.physicalToVirtual(address);
        mapPage(kernel_pml4, virtual, address, flags);
        invalidate(virtual);
    }
    return boot_handoff.physicalToVirtual(physical);
}

/// Like `descend`, but also sets the U/S bit on the intermediate entry (new or
/// pre-existing): ring-3 access requires U at *every* level, and `descend` leaves
/// existing entries untouched. Only used under user-exclusive virtual ranges, so
/// no kernel mapping's protection is widened (the leaf still governs).
fn descendUser(entry: *u64) u64 {
    const table = descend(entry);
    entry.* |= user;
    return table;
}

/// Map one 4 KiB page `virtual` -> `physical` accessible from ring 3. W^X is the
/// caller's contract: code pages are read-only + executable, data pages are
/// writable + no-execute. `virtual` must lie in a user-exclusive region (see
/// `descendUser`).
pub fn mapUser(virtual: u64, physical: u64, writable_page: bool, executable: bool) void {
    mapUserInto(kernel_pml4, virtual, physical, writable_page, executable);
}

/// Map a ring-3-accessible page into the address space rooted at `pml4` (which
/// may be a process's own table or the kernel's). W^X is the caller's contract.
pub fn mapUserInto(pml4: u64, virtual: u64, physical: u64, writable_page: bool, executable: bool) void {
    var flags: u64 = present | user;
    if (writable_page) flags |= writable;
    if (!executable) flags |= no_execute;
    const pml4e = &tableAt(pml4)[(virtual >> 39) & 0x1FF];
    const pdpt = descendUser(pml4e);
    const pdpte = &tableAt(pdpt)[(virtual >> 30) & 0x1FF];
    const pd = descendUser(pdpte);
    const pde = &tableAt(pd)[(virtual >> 21) & 0x1FF];
    const pt = descendUser(pde);
    tableAt(pt)[(virtual >> 12) & 0x1FF] = (physical & address_mask) | flags;
    invalidate(virtual);
}

/// Map a device MMIO window `[physical, physical+len)` into the user (low) half of the
/// address space rooted at `pml4`, page by page. Unlike `mapUserInto` these pages
/// are **strong-uncacheable** (PCD|PWT — device registers must not be cached) and
/// carry the `device_grant` bit so teardown does not return the MMIO frames to the
/// RAM allocator (`freeSubtree`). RW + NX; the caller places `virtual` in a
/// user-exclusive range (PML4[225]). Both `virtual` and `physical` are page-aligned by
/// the caller; a sub-page `physical` offset is the caller's to re-apply.
pub fn mapUserDeviceInto(pml4: u64, virtual: u64, physical: u64, len: u64) void {
    const flags: u64 = present | user | writable | no_execute | pcd | pwt | device_grant;
    const first = physical & ~@as(u64, page_size - 1);
    const last = (physical + (if (len == 0) 1 else len) - 1) & ~@as(u64, page_size - 1);
    var off: u64 = 0;
    while (first + off <= last) : (off += page_size) {
        const v = virtual + off;
        const pml4e = &tableAt(pml4)[(v >> 39) & 0x1FF];
        const pdpt = descendUser(pml4e);
        const pdpte = &tableAt(pdpt)[(v >> 30) & 0x1FF];
        const pd = descendUser(pdpte);
        const pde = &tableAt(pd)[(v >> 21) & 0x1FF];
        const pt = descendUser(pde);
        tableAt(pt)[(v >> 12) & 0x1FF] = ((first + off) & address_mask) | flags;
        invalidate(v);
    }
}

/// Map `[physical, physical+len)` into the user half rooted at `pml4` as **coherent
/// DMA memory**: strong-uncacheable (PCD|PWT — a device reads/writes this RAM without
/// snooping the CPU caches) but, unlike `mapUserDeviceInto`, **without** `device_grant`
/// — because these frames are real RAM from `pmm.allocContiguous`, so teardown
/// (`freeSubtree`) must return them to the allocator like any other user page. RW + NX.
/// The caller aligns `virtual`/`physical` and places `virtual` in the DMA arena.
pub fn mapUserDmaInto(pml4: u64, virtual: u64, physical: u64, len: u64) void {
    const flags: u64 = present | user | writable | no_execute | pcd | pwt;
    const first = physical & ~@as(u64, page_size - 1);
    const last = (physical + (if (len == 0) 1 else len) - 1) & ~@as(u64, page_size - 1);
    var off: u64 = 0;
    while (first + off <= last) : (off += page_size) {
        const v = virtual + off;
        const pml4e = &tableAt(pml4)[(v >> 39) & 0x1FF];
        const pdpt = descendUser(pml4e);
        const pdpte = &tableAt(pdpt)[(v >> 30) & 0x1FF];
        const pd = descendUser(pdpte);
        const pde = &tableAt(pd)[(v >> 21) & 0x1FF];
        const pt = descendUser(pde);
        tableAt(pt)[(v >> 12) & 0x1FF] = ((first + off) & address_mask) | flags;
        invalidate(v);
    }
}

/// Create a new address space: a fresh PML4 with an empty user half and the
/// kernel's higher half shared in (copying PML4[256..512), whose entries point
/// at the kernel's PDPTs — pre-created at init and never restaled, so growth in
/// the kernel half propagates to every address space). Returns the physical
/// PML4, or null if out of frames.
pub fn createAddressSpace() ?u64 {
    const pml4 = alloc_frame() orelse return null;
    const t = tableAt(pml4);
    @memset(t[0..256], 0); // empty user half
    @memcpy(t[256..512], tableAt(kernel_pml4)[256..512]); // shared kernel half
    return pml4;
}

/// Tear down an address space created by `createAddressSpace`: free every frame
/// and table in the user half [0..256), then the PML4 itself. The shared kernel
/// half [256..512) is never touched. The caller must not be running on `pml4`.
pub fn destroyAddressSpace(pml4: u64) void {
    const t = tableAt(pml4);
    for (0..256) |i| {
        if (t[i] & present != 0) freeSubtree(t[i] & address_mask, 3); // PDPT level
    }
    free_frame(pml4);
}

/// Recursively free a page-table subtree: `level` 3 = PDPT, 2 = PD, 1 = PT. At
/// level 1 the entries are leaf data frames; above, they are child tables.
fn freeSubtree(physical: u64, level: u32) void {
    const t = tableAt(physical);
    for (t) |e| {
        if (e & present == 0) continue;
        if (level > 1) {
            freeSubtree(e & address_mask, level - 1);
        } else if (e & device_grant == 0) {
            // A device-grant leaf points at MMIO, not RAM — returning it to the
            // frame allocator would corrupt the pool. Only reclaim real RAM.
            free_frame(e & address_mask);
        }
    }
    free_frame(physical); // page-table frames are always real RAM
}

/// Whether `virtual` is currently mapped **executable** — present with the NX bit
/// clear. Walks the 4-level tables (all danos mappings are 4 KiB, so no huge-page
/// case). Returns false if unmapped. Used for W^X checks in tests.
pub fn isExecutable(virtual: u64) bool {
    const pml4e = tableAt(kernel_pml4)[(virtual >> 39) & 0x1FF];
    if (pml4e & present == 0) return false;
    const pdpte = tableAt(pml4e & address_mask)[(virtual >> 30) & 0x1FF];
    if (pdpte & present == 0) return false;
    const pde = tableAt(pdpte & address_mask)[(virtual >> 21) & 0x1FF];
    if (pde & present == 0) return false;
    const pte = tableAt(pde & address_mask)[(virtual >> 12) & 0x1FF];
    if (pte & present == 0) return false;
    return pte & no_execute == 0;
}

/// Make an already-identity-mapped RAM page **executable** (clear its NX bit),
/// leaving it present and writable. The blanket RAM mapping is NX for W^X, but the
/// application processors fetch the AP trampoline from a low RAM page under paging —
/// so that one page must be executable. A deliberate, temporary W^X exception for a
/// single bring-up page; the caller frees it once every AP is up.
pub fn setExecutable(physical: u64) void {
    mapPage(kernel_pml4, physical, physical, present | writable); // note: no no_execute
    invalidate(physical);
}

/// Remove a mapping and flush it from the TLB.
pub fn unmap(virtual: u64) void {
    unmapInto(kernel_pml4, virtual);
}

/// Remove a mapping from the address space rooted at `pml4` (a process's own
/// table or the kernel's) and flush it from the TLB. Clears only the leaf PTE —
/// the intermediate tables and any frame the PTE pointed at are left to the
/// caller (munmap frees the frame; `destroyAddressSpace` reclaims the tables).
pub fn unmapInto(pml4: u64, virtual: u64) void {
    const pml4e = tableAt(pml4)[(virtual >> 39) & 0x1FF];
    if (pml4e & present == 0) return;
    const pdpte = tableAt(pml4e & address_mask)[(virtual >> 30) & 0x1FF];
    if (pdpte & present == 0) return;
    const pde = tableAt(pdpte & address_mask)[(virtual >> 21) & 0x1FF];
    if (pde & present == 0) return;
    tableAt(pde & address_mask)[(virtual >> 12) & 0x1FF] = 0;
    invalidate(virtual);
}

/// Resolve a virtual address to a physical one in the address space rooted at
/// `pml4`, walking the tables through the physmap (CR3-independent — works for
/// any address space, not just the live one). Returns null if `virtual` is not
/// mapped at any level. All danos mappings are 4 KiB, so there is no huge-page
/// case. The foundation for cross-address-space copies and for munmap (which
/// needs the frame behind a user vaddr to free it).
pub fn translateIn(pml4: u64, virtual: u64) ?u64 {
    const pml4e = tableAt(pml4)[(virtual >> 39) & 0x1FF];
    if (pml4e & present == 0) return null;
    const pdpte = tableAt(pml4e & address_mask)[(virtual >> 30) & 0x1FF];
    if (pdpte & present == 0) return null;
    const pde = tableAt(pdpte & address_mask)[(virtual >> 21) & 0x1FF];
    if (pde & present == 0) return null;
    const pte = tableAt(pde & address_mask)[(virtual >> 12) & 0x1FF];
    if (pte & present == 0) return null;
    return (pte & address_mask) | (virtual & (page_size - 1));
}

fn invalidate(virtual: u64) void {
    // invlpg needs its operand via a register-indirect memory reference that Zig
    // inline asm won't form directly, so stage the address in a register first.
    asm volatile (
        \\mov %[v], %%rax
        \\invlpg (%%rax)
        :
        : [v] "r" (virtual),
        : .{ .rax = true, .memory = true });
}