// ---------------------------------------------------------------------------
// Identity page-table helpers — 4-level paging (x86_64).
//
// The kernel keeps an identity-mapped 4-level tree (no higher half).  At
// boot it runs on UEFI's tables; `adopt_tables()` rebuilds that map into
// our own tables (UEFI maps its own table pages read-only, so the firmware
// tree cannot be modified in place) and switches CR3.  After that the
// helpers walk our tree to map / unmap / query individual 4 KB pages 1:1.
// Page-table pages come from the physical page allocator.  Firmware may
// have used 2 MB / 1 GB large pages for RAM, so the walk splits large pages
// into 4 KB tables on demand.  Split tables stay allocated afterwards (no
// refcounting) — fine for a kernel that owns all memory.
// ---------------------------------------------------------------------------

use crate::memory::{PAGE_SIZE, alloc_pages, free_pages};

pub const PAGE_PRESENT: u64 = 1 << 0;
pub const PAGE_WRITABLE: u64 = 1 << 1;
pub const PAGE_USER: u64 = 1 << 2;
pub const PAGE_PS: u64 = 1 << 7; // large page (2 MB PDE, 1 GB PDPTE)
pub const PAGE_NX: u64 = 1 << 63;

/// Physical address bits in a page-table entry.
const PHYS_ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;

const PAGE_2M: usize = 1 << 21;
const PAGE_1G: usize = 1 << 30;

/// Current PML4 root (physical address, identity-mapped).
pub fn cr3() -> usize {
    let val: usize;
    unsafe {
        core::arch::asm!("mov {0}, cr3", out(reg) val, options(nostack, preserves_flags));
    }
    val
}

/// Full TLB flush by reloading CR3 (slow but bulletproof; map/unmap are
/// rare operations in this kernel).
fn flush_tlb_all() {
    unsafe {
        let cr3: usize;
        core::arch::asm!("mov {0}, cr3", out(reg) cr3, options(nostack, preserves_flags));
        core::arch::asm!("mov cr3, {0}", in(reg) cr3, options(nostack, preserves_flags));
    }
}

/// Rebuild the identity map into our own page tables and switch CR3.
///
/// UEFI maps its own page-table pages read-only (R/W clear), so we cannot
/// modify the firmware tree in place: map/unmap would #PF.  Instead, walk
/// the current tables, replicate every present leaf (4 KB / 2 MB / 1 GB)
/// into fresh tables from the page allocator, force R/W so the tree stays
/// editable, then load the new PML4 into CR3.
pub fn adopt_tables() -> bool {
    unsafe {
        let old_pml4 = cr3() as *const u64;
        let Some(new_pml4) = alloc_pages(1) else {
            return false;
        };
        core::ptr::write_bytes(new_pml4 as *mut u8, 0, PAGE_SIZE);
        let new_pml4 = new_pml4 as *mut u64;

        for l4 in 0..512 {
            let l4e = old_pml4.add(l4).read_volatile();
            if l4e & PAGE_PRESENT == 0 {
                continue;
            }
            // Replicate this PDPT into a fresh one.
            let old_pdpt = (l4e & PHYS_ADDR_MASK) as *const u64;
            let Some(new_pdpt) = alloc_pages(1) else {
                return false;
            };
            core::ptr::write_bytes(new_pdpt as *mut u8, 0, PAGE_SIZE);
            let new_pdpt = new_pdpt as *mut u64;
            *new_pml4.add(l4) =
                (new_pdpt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (l4e & PAGE_USER);

            for l3 in 0..512 {
                let pdpte = old_pdpt.add(l3).read_volatile();
                if pdpte & PAGE_PRESENT == 0 {
                    continue;
                }
                if pdpte & PAGE_PS != 0 {
                    // 1 GB large page — copy the leaf, force writable.
                    *new_pdpt.add(l3) = pdpte | PAGE_WRITABLE;
                    continue;
                }
                let old_pd = (pdpte & PHYS_ADDR_MASK) as *const u64;
                let Some(new_pd) = alloc_pages(1) else {
                    return false;
                };
                core::ptr::write_bytes(new_pd as *mut u8, 0, PAGE_SIZE);
                let new_pd = new_pd as *mut u64;
                *new_pdpt.add(l3) =
                    (new_pd as u64) | PAGE_PRESENT | PAGE_WRITABLE | (pdpte & PAGE_USER);

                for l2 in 0..512 {
                    let pde = old_pd.add(l2).read_volatile();
                    if pde & PAGE_PRESENT == 0 {
                        continue;
                    }
                    if pde & PAGE_PS != 0 {
                        // 2 MB large page — copy the leaf, force writable.
                        *new_pd.add(l2) = pde | PAGE_WRITABLE;
                        continue;
                    }
                    let old_pt = (pde & PHYS_ADDR_MASK) as *const u64;
                    let Some(new_pt) = alloc_pages(1) else {
                        return false;
                    };
                    core::ptr::write_bytes(new_pt as *mut u8, 0, PAGE_SIZE);
                    let new_pt = new_pt as *mut u64;
                    *new_pd.add(l2) =
                        (new_pt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (pde & PAGE_USER);

                    for l1 in 0..512 {
                        let pte = old_pt.add(l1).read_volatile();
                        if pte & PAGE_PRESENT != 0 {
                            *new_pt.add(l1) = pte | PAGE_WRITABLE;
                        }
                    }
                }
            }
        }

        // Switch to the new tables (CR3 reload flushes the TLB).
        let new_cr3 = new_pml4 as usize;
        core::arch::asm!("mov cr3, {0}", in(reg) new_cr3, options(nostack, preserves_flags));
        true
    }
}

/// Ensure the PDPT entry points at a page-directory table, splitting a
/// 1 GB large page at `base_1g` if present.  Returns the PD pointer.
unsafe fn ensure_pd(pdpte: &mut u64, base_1g: usize) -> Option<*mut u64> {
    unsafe {
        if *pdpte & PAGE_PRESENT == 0 {
            let page = alloc_pages(1)? as *mut u64;
            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
            *pdpte = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
            return Some(page);
        }
        if *pdpte & PAGE_PS != 0 {
            // Split the 1 GB page into 512 × 2 MB pages.
            let pd = alloc_pages(1)? as *mut u64;
            core::ptr::write_bytes(pd as *mut u8, 0, PAGE_SIZE);
            let inherit = *pdpte & (PAGE_USER | PAGE_NX);
            for i in 0..512 {
                let p = base_1g + i * PAGE_2M;
                *pd.add(i) = (p as u64) | PAGE_PRESENT | PAGE_WRITABLE | PAGE_PS | inherit;
            }
            *pdpte = (pd as u64) | PAGE_PRESENT | PAGE_WRITABLE | (*pdpte & PAGE_USER);
            return Some(pd);
        }
        Some(((*pdpte & PHYS_ADDR_MASK) as usize) as *mut u64)
    }
}

/// Ensure the PD entry points at a page table, splitting a 2 MB large page
/// at `base_2m` if present.  Returns the PT pointer.
unsafe fn ensure_pt(pde: &mut u64, base_2m: usize) -> Option<*mut u64> {
    unsafe {
        if *pde & PAGE_PRESENT == 0 {
            let page = alloc_pages(1)? as *mut u64;
            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
            *pde = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
            return Some(page);
        }
        if *pde & PAGE_PS != 0 {
            // Split the 2 MB page into 512 × 4 KB pages.
            let pt = alloc_pages(1)? as *mut u64;
            core::ptr::write_bytes(pt as *mut u8, 0, PAGE_SIZE);
            let inherit = *pde & (PAGE_USER | PAGE_NX);
            for i in 0..512 {
                let p = base_2m + i * PAGE_SIZE;
                *pt.add(i) = (p as u64) | PAGE_PRESENT | PAGE_WRITABLE | inherit;
            }
            *pde = (pt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (*pde & PAGE_USER);
            return Some(pt);
        }
        Some(((*pde & PHYS_ADDR_MASK) as usize) as *mut u64)
    }
}

/// Identity-map the physical page at `phys` (must be 4 KB aligned).
/// `flags` adds to PRESENT (e.g. PAGE_WRITABLE, PAGE_USER).
pub fn map_page(phys: usize, flags: u64) -> bool {
    if !phys.is_multiple_of(PAGE_SIZE) {
        return false;
    }
    unsafe {
        let pml4 = cr3() as *mut u64;
        let l4 = (phys >> 39) & 0x1FF;
        let l3 = (phys >> 30) & 0x1FF;
        let l2 = (phys >> 21) & 0x1FF;
        let l1 = (phys >> 12) & 0x1FF;

        let pml4e = &mut *pml4.add(l4);
        if *pml4e & PAGE_PRESENT == 0 {
            let Some(page) = alloc_pages(1) else {
                return false;
            };
            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
            *pml4e = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
        }
        // PML4 entries are never large pages.

        let base_1g = phys & !(PAGE_1G - 1);
        let pdpte = &mut *(((*pml4e & PHYS_ADDR_MASK) as usize) as *mut u64).add(l3);
        let Some(pd) = ensure_pd(pdpte, base_1g) else {
            return false;
        };

        let base_2m = phys & !(PAGE_2M - 1);
        let pde = &mut *pd.add(l2);
        let Some(pt) = ensure_pt(pde, base_2m) else {
            return false;
        };

        let pte = &mut *pt.add(l1);
        *pte = (phys as u64) | flags | PAGE_PRESENT;
        flush_tlb_all();
        true
    }
}

/// Identity-map `count` consecutive pages starting at `phys`.
pub fn map_range(phys: usize, count: usize, flags: u64) -> bool {
    for i in 0..count {
        if !map_page(phys + i * PAGE_SIZE, flags) {
            return false;
        }
    }
    true
}

/// Unmap the 4 KB page at `phys`.  Splits large pages as needed so only
/// this one page is removed.  Returns false if it wasn't mapped.
pub fn unmap_page(phys: usize) -> bool {
    if !phys.is_multiple_of(PAGE_SIZE) {
        return false;
    }
    unsafe {
        let pml4 = cr3() as *mut u64;
        let l4 = (phys >> 39) & 0x1FF;
        let l3 = (phys >> 30) & 0x1FF;
        let l2 = (phys >> 21) & 0x1FF;
        let l1 = (phys >> 12) & 0x1FF;

        let pml4e = &mut *pml4.add(l4);
        if *pml4e & PAGE_PRESENT == 0 {
            return false;
        }
        let base_1g = phys & !(PAGE_1G - 1);
        let pdpte = &mut *(((*pml4e & PHYS_ADDR_MASK) as usize) as *mut u64).add(l3);
        if *pdpte & PAGE_PRESENT == 0 {
            return false;
        }
        let Some(pd) = ensure_pd(pdpte, base_1g) else {
            return false;
        };
        let base_2m = phys & !(PAGE_2M - 1);
        let pde = &mut *pd.add(l2);
        if *pde & PAGE_PRESENT == 0 {
            return false;
        }
        let Some(pt) = ensure_pt(pde, base_2m) else {
            return false;
        };
        let pte = &mut *pt.add(l1);
        if *pte & PAGE_PRESENT == 0 {
            return false;
        }
        *pte = 0;
        flush_tlb_all();
        true
    }
}

/// Is the page at `phys` mapped (present at any level, incl. large pages)?
pub fn is_mapped(phys: usize) -> bool {
    unsafe {
        let pml4 = cr3() as *const u64;
        let l4 = (phys >> 39) & 0x1FF;
        let l3 = (phys >> 30) & 0x1FF;
        let l2 = (phys >> 21) & 0x1FF;
        let l1 = (phys >> 12) & 0x1FF;

        let pml4e = *pml4.add(l4);
        if pml4e & PAGE_PRESENT == 0 {
            return false;
        }
        let pdpte = *((pml4e & PHYS_ADDR_MASK) as *const u64).add(l3);
        if pdpte & PAGE_PRESENT == 0 {
            return false;
        }
        if pdpte & PAGE_PS != 0 {
            return true; // 1 GB page covers phys
        }
        let pde = *((pdpte & PHYS_ADDR_MASK) as *const u64).add(l2);
        if pde & PAGE_PRESENT == 0 {
            return false;
        }
        if pde & PAGE_PS != 0 {
            return true; // 2 MB page covers phys
        }
        let pte = *((pde & PHYS_ADDR_MASK) as *const u64).add(l1);
        pte & PAGE_PRESENT != 0
    }
}

/// Smoke test: allocate two physical pages, unmap one, verify the mapping
/// state, remap it, verify data still works, and free both.
pub fn self_test() -> bool {
    let Some(p1) = alloc_pages(1) else {
        return false;
    };
    let Some(p2) = alloc_pages(1) else {
        return false;
    };

    // Both pages must already be identity-mapped (they're RAM).
    if !is_mapped(p1) || !is_mapped(p2) {
        return false;
    }

    // Unmap p1 — only p1 should disappear.
    if !unmap_page(p1) {
        return false;
    }
    if is_mapped(p1) {
        return false;
    }
    if !is_mapped(p2) {
        return false;
    }

    // Remap and verify we can write through the mapping.
    if !map_page(p1, PAGE_WRITABLE) {
        return false;
    }
    if !is_mapped(p1) {
        return false;
    }
    unsafe {
        core::ptr::write_volatile(p1 as *mut u8, 0xCD);
        if core::ptr::read_volatile(p1 as *const u8) != 0xCD {
            return false;
        }
    }

    free_pages(p1, 1);
    free_pages(p2, 1);
    true
}