use crate::uefi::*;
use core::alloc::{GlobalAlloc, Layout};
use core::fmt;
use core::mem;
use core::ptr;

const MAX_MAP_ENTRIES: usize = 512;

static mut MEMORY_MAP: [MemoryDescriptor; MAX_MAP_ENTRIES] = unsafe { mem::zeroed() };
static mut MAP_ENTRIES: usize = 0;

/// The map key returned by GetMemoryMap, needed for ExitBootServices.
static mut MAP_KEY: usize = 0;

/// Static working buffer for GetMemoryMap (64 KB, way more than enough).
/// Lives in .bss; no pool allocation, so no map-key staleness.
static mut MAP_WORK_BUF: [u8; 65536] = [0u8; 65536];

// ---------------------------------------------------------------------------
// Free-list heap allocator — first fit, splitting, coalescing.
// Used as Rust's #[global_allocator] once the heap region is set up.
//
// Every block (free or allocated) starts with a 16-byte header:
//
//     [ size: usize ][ free: usize ]
//
//   - `size` covers the whole block including the header.
//   - `free` is 1 for free blocks, 0 for allocated blocks.
//
// Free blocks are linked into one address-sorted doubly-linked list; the
// prev/next links live in the block's payload (free space is unused, so
// it's free real estate for the links).  Allocating walks the list for the
// first fit, splits the remainder into a fresh free block, and records the
// block address in an 8-byte back-pointer immediately before the returned
// payload — needed because alignment can leave a gap between the block
// header and the payload.  Freeing reads the back-pointer to find its
// block, marks it free, re-inserts it in address order, and coalesces with
// any adjacent free neighbours.
//
// All sizes are rounded to 16 bytes so block boundaries stay 16-aligned
// and every payload is at least 16-byte aligned.
// ---------------------------------------------------------------------------

const HEADER_SIZE: usize = 16; // size + free flag
const MINI_SIZE: usize = 8; // back-pointer before the payload
const MIN_BLOCK: usize = 32; // smallest block we will create (header + links)
const MIN_ALIGN: usize = 16;

/// Block header — first 16 bytes of every block.
#[repr(C)]
struct Header {
    size: usize,
    free: usize,
}

impl Header {
    /// Where the free-list links live (only meaningful while free).
    #[inline]
    fn links(&self) -> *mut FreeLinks {
        unsafe { (self as *const Header as *mut u8).add(HEADER_SIZE) as *mut FreeLinks }
    }
}

/// Free-list links, stored in the payload of free blocks only.
#[repr(C)]
struct FreeLinks {
    prev: *mut Header,
    next: *mut Header,
}

/// Head of the address-sorted free list (lowest address first).
static mut FREE_LIST: *mut Header = ptr::null_mut();

/// TEST HOOK: fill every free block's payload with 0xAA (header + links
/// preserved).  Mirrors the OptiPlex, whose heap pages carry 0xAA
/// firmware leftover — any code that reads uninitialized heap memory
/// then sees the same pattern the real machine shows in its crashes
/// (ptr=0xaaaaaa…, allocation-size panics, zeroed block headers).
pub fn poison_free_blocks_aa() {
    unsafe {
        let mut b = FREE_LIST;
        let mut blocks = 0u32;
        let mut bytes = 0u64;
        while !b.is_null() {
            let size = (*b).size;
            let off = HEADER_SIZE + core::mem::size_of::<FreeLinks>();
            if size > off {
                let p = (b as *mut u8).add(off);
                let n = size - off;
                // 8-byte words, then a byte tail (blocks are 16-aligned).
                let mut i = 0;
                while i + 8 <= n {
                    (p.add(i) as *mut u64).write_volatile(0xAAAAAAAAAAAAAAAA);
                    i += 8;
                }
                while i < n {
                    p.add(i).write_volatile(0xAA);
                    i += 1;
                }
                bytes += n as u64;
            }
            blocks += 1;
            b = (*(*b).links()).next;
        }
        crate::serial_println!(
            "heap: poisoned {} free block(s), {} KiB with 0x00AAAAAA",
            blocks,
            bytes / 1024
        );
    }
}

/// TEST HOOK: fill every free page (per the bitmap) with the OptiPlex's
/// 0x00AAAAAA scrub pattern — the firmware's leftover fills pages the
/// page allocator will hand out (task stacks, heap growth), so any
/// uninitialized read there matches the real machine byte-for-byte.
pub fn poison_free_pages_aa() {
    unsafe {
        if BITMAP.is_null() {
            crate::serial_println!("pages: poison skipped (no bitmap)");
            return;
        }
        let mut n = 0u32;
        for idx in 0..SPAN_PAGES {
            let free = (*BITMAP.add(idx / 8) >> (idx % 8)) & 1 == 0;
            if !free {
                continue;
            }
            let base = (BASE_PAGE + idx) * PAGE_SIZE;
            let p = base as *mut u64;
            for k in 0..(PAGE_SIZE / 8) {
                p.add(k).write_volatile(0x00AAAAAA00AAAAAA);
            }
            n += 1;
        }
        crate::serial_println!("pages: poisoned {} free page(s) with 0x00AAAAAA", n);
    }
}

/// Heap region bounds (set by `init_heap`).
static mut HEAP_START: usize = 0;
static mut HEAP_END: usize = 0;
/// Total heap capacity — the original region plus pages added by growth.
static mut HEAP_CAPACITY: usize = 0;

/// Kernel image bounds (base..end) captured from the UEFI LoadedImage
/// protocol before ExitBootServices.  Both the heap free list and the
/// page-allocator bitmap reserve this range so no allocation can ever
/// land on the running kernel's own code/data/bss pages.
static mut KERNEL_IMAGE_START: usize = 0;
static mut KERNEL_IMAGE_END: usize = 0;

/// Record the loaded-image memory range (from EFI_LOADED_IMAGE_PROTOCOL).
pub fn set_kernel_image_range(base: usize, size: usize) {
    unsafe {
        KERNEL_IMAGE_START = base;
        KERNEL_IMAGE_END = base + size;
    }
}

/// The loaded kernel image range `(start, end)`, if known.  Used to
/// resolve fault RIPs against the symbol table and to filter stack
/// scans for return addresses.
pub fn kernel_image_range() -> Option<(usize, usize)> {
    unsafe {
        if KERNEL_IMAGE_END > KERNEL_IMAGE_START {
            Some((KERNEL_IMAGE_START, KERNEL_IMAGE_END))
        } else {
            None
        }
    }
}

/// The in-memory `.text` section range `(start, end)`, parsed from the
/// loaded PE headers.  Used to filter fault backtraces: only code
/// addresses are meaningful return addresses (the rest of the image
/// range is .rdata/.data pointers).  Falls back to the full image
/// range if the headers are unreadable.
pub fn text_range() -> Option<(usize, usize)> {
    let (base, _end) = kernel_image_range()?;
    unsafe {
        let base = base as *const u8;
        if base.read() != b'M' || base.add(1).read() != b'Z' {
            return kernel_image_range();
        }
        let e_lfanew = (base.add(0x3c) as *const u32).read() as usize;
        let pe = base.add(e_lfanew);
        if (pe as *const u32).read() != 0x4550 {
            return kernel_image_range();
        }
        let fh = pe.add(4); // COFF file header (20 bytes)
        let size_opt = (fh.add(16) as *const u16).read() as usize;
        let nsec = (fh.add(2) as *const u16).read() as usize;
        let sec = fh.add(20).add(size_opt);
        for i in 0..nsec {
            let s = sec.add(i * 40);
            let nm = core::slice::from_raw_parts(s, 8);
            if nm.starts_with(b".text") {
                let vsz = (s.add(8) as *const u32).read() as usize;
                let va = (s.add(12) as *const u32).read() as usize;
                return Some((base as usize + va, base as usize + va + vsz));
            }
        }
        kernel_image_range()
    }
}

/// Minimum pages `grow_heap` takes from the page allocator per growth.
const HEAP_GROW_PAGES: usize = 64;

#[inline]
fn align_up(x: usize, a: usize) -> usize {
    (x + a - 1) & !(a - 1)
}

/// Insert `block` into the address-sorted free list.
/// Caller guarantees `block` is not already in the list.
unsafe fn free_list_insert(block: *mut Header) {
    let addr = block as usize;
    let mut prev: *mut Header = ptr::null_mut();
    let mut cur = unsafe { FREE_LIST };
    while !cur.is_null() && (cur as usize) < addr {
        prev = cur;
        cur = unsafe { (*(*cur).links()).next };
    }
    unsafe {
        let links = (*block).links();
        (*links).prev = prev;
        (*links).next = cur;
        if prev.is_null() {
            FREE_LIST = block;
        } else {
            (*(*prev).links()).next = block;
        }
        if !cur.is_null() {
            (*(*cur).links()).prev = block;
        }
    }
}

/// Unlink `block` from the free list.
unsafe fn free_list_remove(block: *mut Header) {
    unsafe {
        let links = (*block).links();
        let prev = (*links).prev;
        let next = (*links).next;
        if prev.is_null() {
            FREE_LIST = next;
        } else {
            (*(*prev).links()).next = next;
        }
        if !next.is_null() {
            (*(*next).links()).prev = prev;
        }
    }
}

/// Merge `block` with any adjacent free neighbours (both directions).
/// `block` must already be in the free list.
unsafe fn coalesce(block: *mut Header) {
    unsafe {
        // Absorb the next free block if it sits right after us.
        let next = (*(*block).links()).next;
        if !next.is_null() && (block as usize) + (*block).size == next as usize {
            (*block).size += (*next).size;
            let next_links = (*next).links();
            (*(*block).links()).next = (*next_links).next;
            if !(*next_links).next.is_null() {
                (*(*(*next_links).next).links()).prev = block;
            }
        }
        // Be absorbed by the previous free block if it ends right before us.
        let prev = (*(*block).links()).prev;
        if !prev.is_null() && (prev as usize) + (*prev).size == block as usize {
            (*prev).size += (*block).size;
            let block_links = (*block).links();
            (*(*prev).links()).next = (*block_links).next;
            if !(*block_links).next.is_null() {
                (*(*(*block_links).next).links()).prev = prev;
            }
        }
    }
}

/// First-fit search: find the first block big enough for `size` aligned to
/// `align`, split it, and return the aligned payload (back-pointer written).
/// Returns null when no block fits.
unsafe fn alloc_from_free_list(size: usize, align: usize) -> *mut u8 {
    unsafe {
        let mut cur = FREE_LIST;
        while !cur.is_null() {
            let block = cur;
            let block_size = (*block).size;
            // Payload sits after the header + back-pointer, aligned up.
            let payload = align_up(block as usize + HEADER_SIZE + MINI_SIZE, align);
            if payload + size <= (block as usize) + block_size {
                let consumed = payload + size - (block as usize);
                let remainder = block_size - consumed;

                free_list_remove(block);

                if remainder >= MIN_BLOCK {
                    // Split: carve the tail into a fresh free block.
                    let rem = (payload + size) as *mut Header;
                    (*rem).size = remainder;
                    (*rem).free = 1;
                    free_list_insert(rem);
                    (*block).size = consumed;
                }
                (*block).free = 0;

                // Back-pointer so `free` can find the block even when
                // alignment left a gap before the payload.
                *((payload as *mut usize).sub(1)) = block as usize;
                return payload as *mut u8;
            }
            cur = (*(*block).links()).next;
        }
        ptr::null_mut()
    }
}

/// Expand the heap by taking `pages` contiguous physical pages from the
/// page allocator and adding them to the free list as one free block.
/// Returns false when the page allocator can't provide them.
unsafe fn grow_heap(pages: usize) -> bool {
    unsafe {
        if let Some(addr) = alloc_pages(pages) {
            let block = addr as *mut Header;
            (*block).size = pages * PAGE_SIZE;
            (*block).free = 1;
            free_list_insert(block);
            coalesce(block);
            HEAP_CAPACITY += pages * PAGE_SIZE;
            true
        } else {
            false
        }
    }
}

/// The kernel heap allocator.
struct Allocator;

unsafe impl Sync for Allocator {}

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let size = align_up(layout.size().max(1), MIN_ALIGN);
        let align = layout.align().max(MIN_ALIGN);
        unsafe {
            // First fit; when the free list is exhausted, grow the heap
            // from the page allocator and retry.  Growth is proportional
            // to the request (min HEAP_GROW_PAGES), so large allocations
            // converge quickly.
            loop {
                let p = alloc_from_free_list(size, align);
                if !p.is_null() {
                    return p;
                }
                let need = (size / PAGE_SIZE + 1).max(HEAP_GROW_PAGES);
                if !grow_heap(need) {
                    return ptr::null_mut();
                }
            }
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if ptr.is_null() {
            return;
        }
        unsafe {
            // Real-hardware corruption guard: the back-pointer at [ptr-8]
            // must point at a sane block header inside the heap.  A
            // garbage back-pointer (heap/stack corruption clobbered the
            // block header, or the caller passed a bogus pointer) used to
            // turn into a #GP on the store below — now it turns into a
            // readable report + halt instead.
            let p = ptr as usize;
            if !back_pointer_ok(p) {
                let rsp: usize;
                core::arch::asm!(
                    "mov {}, rsp",
                    out(reg) rsp,
                    options(nostack, preserves_flags)
                );
                heap_guard_report(p, rsp, layout.size(), layout.align());
            }
            let block = *((ptr as *const usize).sub(1)) as *mut Header;
            debug_assert!(
                (*block).free == 0,
                "heap: double free at {:#x}",
                ptr as usize
            );
            (*block).free = 1;
            free_list_insert(block);
            coalesce(block);
        }
    }
}

#[global_allocator]
static ALLOCATOR: Allocator = Allocator;

/// Current kernel-heap extent, if the heap has been initialised.
pub(crate) fn heap_range() -> Option<(usize, usize)> {
    unsafe {
        if HEAP_END > HEAP_START {
            Some((HEAP_START, HEAP_END))
        } else {
            None
        }
    }
}

/// True when `[ptr-8]` holds a plausible back-pointer: the *value* there
/// must be a block address inside the heap whose header size is sane and
/// whose free flag says the block is still allocated (a block being freed
/// must be in-use).
unsafe fn back_pointer_ok(ptr: usize) -> bool {
    unsafe {
        let Some((hs, he)) = heap_range() else {
            return true; // heap not initialised yet — trust the pointer
        };
        let slot = ptr.wrapping_sub(8);
        if slot < hs || slot >= he {
            // The slot itself is outside the heap — can't even read it
            // safely (a garbage ptr reads random memory).
            return false;
        }
        let block = *(slot as *const usize);
        if block < hs || block >= he {
            return false; // back-pointer points outside the heap
        }
        let size = *(block as *const usize);
        let free = *((block + 8) as *const usize);
        size >= MIN_BLOCK && block + size <= he && free == 0
    }
}

/// Heap-corruption report, called from `dealloc` when the back-pointer
/// check fails.  Heap-free by construction (we ARE the allocator): raw
/// serial, the crash band, crash.log, then halt.
fn heap_guard_report(ptr: usize, rsp: usize, size: usize, align: usize) -> ! {
    let (hs, he) = heap_range().unwrap_or((0, 0));
    let block = if ptr >= hs + 8 && ptr <= he { unsafe { *(ptr.wrapping_sub(8) as *const usize) } } else { 0 };

    let mut b1 = [0u8; 96];
    let n1 = {
        let mut w = crate::strbuf::FmtBuf::new(&mut b1);
        let _ = fmt::write(&mut w, core::format_args!("!!! heap corruption in dealloc"));
        w.len()
    };
    let mut b2 = [0u8; 96];
    let n2 = {
        let mut w = crate::strbuf::FmtBuf::new(&mut b2);
        let _ = fmt::write(&mut w, core::format_args!("ptr={:#x} block={:#x}", ptr, block));
        w.len()
    };
    let mut b3 = [0u8; 96];
    let n3 = {
        let mut w = crate::strbuf::FmtBuf::new(&mut b3);
        let _ = fmt::write(
            &mut w,
            core::format_args!(
                "heap={:#x}..{:#x} rsp={:#x}",
                hs, he, rsp
            ),
        );
        w.len()
    };
    let mut b3b = [0u8; 96];
    let n3b = {
        let mut w = crate::strbuf::FmtBuf::new(&mut b3b);
        let _ = fmt::write(
            &mut w,
            core::format_args!("drop size={:#x} align={:#x}", size, align),
        );
        w.len()
    };

    // Candidate block bases for this ptr (payload sits at block+24 for
    // align-1 allocs, block+32 for align-16) with their size/free
    // fields, plus a longer raw window around ptr so the next crash
    // shows the header region and neighbours, not just the payload.
    let mut b4 = [0u8; 96];
    let mut b4b = [0u8; 96];
    let mut b4c = [0u8; 96];
    let mut b4d = [0u8; 96];
    let focus = if block >= hs && block < he { block } else { ptr };
    let mut a = (focus.wrapping_sub(48)).max(hs) & !7;
    let (n4, n4b, n4c, n4d);
    if focus >= hs + 8 && focus < he {
        let mut w = crate::strbuf::FmtBuf::new(&mut b4);
        let _ = fmt::write(&mut w, core::format_args!("mem:"));
        for _ in 0..4 {
            let _ = fmt::write(
                &mut w,
                core::format_args!(" {:x}={:#x}", a, unsafe { *(a as *const usize) }),
            );
            a += 8;
        }
        n4 = w.len();
        let mut w = crate::strbuf::FmtBuf::new(&mut b4b);
        let _ = fmt::write(&mut w, core::format_args!("mem2:"));
        for _ in 0..4 {
            let _ = fmt::write(
                &mut w,
                core::format_args!(" {:x}={:#x}", a, unsafe { *(a as *const usize) }),
            );
            a += 8;
        }
        n4b = w.len();
        let mut w = crate::strbuf::FmtBuf::new(&mut b4c);
        let _ = fmt::write(&mut w, core::format_args!("mem3:"));
        for _ in 0..4 {
            let _ = fmt::write(
                &mut w,
                core::format_args!(" {:x}={:#x}", a, unsafe { *(a as *const usize) }),
            );
            a += 8;
        }
        n4c = w.len();
        let mut w = crate::strbuf::FmtBuf::new(&mut b4d);
        let _ = fmt::write(&mut w, core::format_args!("mem4:"));
        for _ in 0..4 {
            let _ = fmt::write(
                &mut w,
                core::format_args!(" {:x}={:#x}", a, unsafe { *(a as *const usize) }),
            );
            a += 8;
        }
        n4d = w.len();
    } else {
        let mut w = crate::strbuf::FmtBuf::new(&mut b4);
        let _ = fmt::write(&mut w, core::format_args!("mem: out of heap — no dump"));
        n4 = w.len();
        n4b = 0;
        n4c = 0;
        n4d = 0;
    };
    // Explicit candidate-block header fields (size/free at +0/+8).
    let mut b4e = [0u8; 96];
    let ne = if ptr >= hs + 24 && ptr <= he {
        let hb = ptr - 24;
        let hb2 = ptr - 32;
        let mut w = crate::strbuf::FmtBuf::new(&mut b4e);
        let _ = fmt::write(
            &mut w,
            core::format_args!(
                "hdr: b={:#x} s={:#x} f={:#x} b16={:#x} s={:#x} f={:#x}",
                hb,
                unsafe { *(hb as *const usize) },
                unsafe { *((hb + 8) as *const usize) },
                hb2,
                unsafe { *(hb2 as *const usize) },
                unsafe { *((hb2 + 8) as *const usize) },
            ),
        );
        w.len()
    } else {
        0
    };

    // Live return-address chain on the stack above our rsp (dealloc ←
    // drop glue ← shell::run); raw stack words, resolves against the PDB.
    // 16 words in 4 lines of 4 — enough to capture the drop caller
    // chain on the band even when the copy truncates a line.
    let mut b5 = [0u8; 96];
    let mut b6 = [0u8; 96];
    let mut b7 = [0u8; 96];
    let mut b8 = [0u8; 96];
    let (n5, n6, n7, n8) = {
        let mut a = rsp.wrapping_add(8);
        let bufs: [&mut [u8]; 4] = [&mut b5, &mut b6, &mut b7, &mut b8];
        let mut lens = [0usize; 4];
        for li in 0..4 {
            let mut w = crate::strbuf::FmtBuf::new(bufs[li]);
            let lbl = match li {
                0 => "",
                1 => "2",
                2 => "3",
                _ => "4",
            };
            let _ = fmt::write(
                &mut w,
                core::format_args!("stk{}:", lbl),
            );
            for _ in 0..4 {
                let v = unsafe { *(a as *const usize) };
                let _ = fmt::write(&mut w, core::format_args!(" {:x}={:#x}", a, v));
                a += 8;
            }
            lens[li] = w.len();
        }
        (lens[0], lens[1], lens[2], lens[3])
    };

    let l1 = core::str::from_utf8(&b1[..n1]).unwrap_or("!!!");
    let l2 = core::str::from_utf8(&b2[..n2]).unwrap_or("");
    let l3 = core::str::from_utf8(&b3[..n3]).unwrap_or("");
    let l3b = core::str::from_utf8(&b3b[..n3b]).unwrap_or("");
    let l4 = core::str::from_utf8(&b4[..n4]).unwrap_or("");
    let l4b = if n4b > 0 { core::str::from_utf8(&b4b[..n4b]).unwrap_or("") } else { "" };
    let l4c = if n4c > 0 { core::str::from_utf8(&b4c[..n4c]).unwrap_or("") } else { "" };
    let l4d = if n4d > 0 { core::str::from_utf8(&b4d[..n4d]).unwrap_or("") } else { "" };
    let l4e = if ne > 0 { core::str::from_utf8(&b4e[..ne]).unwrap_or("") } else { "" };
    let l5 = core::str::from_utf8(&b5[..n5]).unwrap_or("");
    let l6 = core::str::from_utf8(&b6[..n6]).unwrap_or("");
    let l7 = core::str::from_utf8(&b7[..n7]).unwrap_or("");
    let l8 = core::str::from_utf8(&b8[..n8]).unwrap_or("");
    for l in [l2, l3, l3b, l4, l4b, l4c, l4d, l4e, l5, l6, l7, l8] {
        crate::serial::write_raw(l);
        crate::serial::write_raw("\r\n");
    }
    let lines = [l1, l2, l3, l3b, l4, l4b, l4c, l4d, l4e, l5, l6, l7, l8];
    crate::gop::dump_lines(&lines);
    crate::fs::write_crash_log(&lines);
    crate::interrupts::halt()
}

/// Capture the UEFI memory map into the static buffer.  Uses a pre-allocated
/// .bss working buffer so no pool allocation (and thus no memory-map change)
/// occurs between GetMemoryMap and ExitBootServices.
pub fn capture_memory_map(bs: &BootServices) -> Result<usize, Status> {
    capture_memory_map_impl(bs, true)
}

/// Re-capture the memory map and return a FRESH ExitBootServices key
/// WITHOUT rebuilding the heap free list.  Used by the EBS retry path:
/// a rebuild from the raw map wipes the allocation state of every
/// pre-EBS live allocation (the display shadow, the vt grids) — their
/// spans then become free while still in use, and later allocations
/// land inside the live shadow (the OptiPlex crash family).
pub fn refresh_map_key(bs: &BootServices) -> Result<usize, Status> {
    capture_memory_map_impl(bs, false)
}

fn capture_memory_map_impl(bs: &BootServices, rebuild: bool) -> Result<usize, Status> {
    let mut map_key: usize = 0;
    let mut desc_size: usize = 0;
    let mut desc_ver: u32 = 0;

    // Probe with null to get required size + descriptor size
    let mut buf_size: usize = 0;
    let status = unsafe {
        (bs.get_memory_map)(
            &mut buf_size,
            ptr::null_mut(),
            &mut map_key,
            &mut desc_size,
            &mut desc_ver,
        )
    };
    if status != EFI_BUFFER_TOO_SMALL {
        // Fallback: allocate a sensible size
        buf_size = 64 * 1024;
    }
    if desc_size == 0 {
        desc_size = mem::size_of::<MemoryDescriptor>();
    }

    // Round up to desc_size alignment
    let buf_size = buf_size.div_ceil(desc_size) * desc_size;

    // If our static buffer is too small, fall back to pool (rare)
    if buf_size > unsafe { (*core::ptr::addr_of!(MAP_WORK_BUF)).len() } {
        return Err(EFI_BUFFER_TOO_SMALL);
    }

    let mut actual_size = buf_size;
    let status = unsafe {
        (bs.get_memory_map)(
            &mut actual_size,
            &raw mut MAP_WORK_BUF[0] as *mut MemoryDescriptor,
            &mut map_key,
            &mut desc_size,
            &mut desc_ver,
        )
    };

    if status == EFI_BUFFER_TOO_SMALL {
        // The map grew between the probe and the real call.  Try a larger
        // static buffer (can't grow here without pool, so return error to
        // let the caller retry if they can allocate pool).
        // For our 64 KB buffer this should never happen on OVMF.
        return Err(EFI_BUFFER_TOO_SMALL);
    }
    if status != EFI_SUCCESS {
        return Err(status);
    }

    // Copy descriptors into our indexed array
    let count = actual_size / desc_size;
    let count = count.min(MAX_MAP_ENTRIES);
    unsafe {
        for i in 0..count {
            let p = (&raw const MAP_WORK_BUF[..]) as *const u8;
            let off = i * desc_size;
            let src = p.add(off) as *const MemoryDescriptor;
            MEMORY_MAP[i] = ptr::read(src);
        }
        MAP_ENTRIES = count;
        MAP_KEY = map_key;
    }

    // Do NOT free anything — the static buffer stays.  Map key is still
    // valid because no pool allocation/free happened after GetMemoryMap.

    if rebuild {
        init_heap();
    }
    Ok(map_key)
}

/// Retry ExitBootServices with a fresh memory map.
///
/// # Safety
///
/// `image_handle` must be the valid EFI image handle this kernel was
/// started with (opaque firmware pointer; never dereferenced here).
pub unsafe fn refresh_and_exit(bs: &BootServices, image_handle: Handle) -> Result<usize, Status> {
    // Fresh key WITHOUT rebuilding the heap free list: a rebuild from the
    // raw map wipes the allocation state of every pre-EBS live allocation
    // (the display shadow, the vt grids) — their spans then become free
    // while still in use, and later allocations land inside the live
    // shadow (the OptiPlex crash family: display writes clobber heap
    // objects, zeroed back-pointers, gray-text pixels).
    let key = refresh_map_key(bs)?;
    let status = unsafe { (bs.exit_boot_services)(image_handle, key) };
    if status == EFI_SUCCESS {
        Ok(key)
    } else {
        Err(status)
    }
}

fn init_heap() {
    unsafe {
        // Regions that must never be handed out as heap memory, even when
        // the firmware lists them as conventional:
        //   * the GOP framebuffer (heap allocations overwriting the display),
        //   * the kernel image itself (a heap allocation landing on the
        //     kernel's .bss would silently corrupt statics — observed as a
        //     garbage ACTIVE_VOL and a dead filesystem in the fault path).
        let holes: [(usize, usize); 2] = [
            crate::gop::framebuffer_region().unwrap_or((0, 0)),
            kernel_image_range().unwrap_or((0, 0)),
        ];
        let mut best_start = 0usize;
        let mut best_size = 0usize;
        for i in 0..MAP_ENTRIES {
            let d = &MEMORY_MAP[i];
            if d.kind == MEMORY_TYPE_CONVENTIONAL {
                let s = d.physical_start as usize;
                let e = s + d.number_of_pages as usize * 4096;
                // Split at each hole so the free list NEVER contains
                // reserved pages: a coarse size subtraction leaves one
                // block *spanning* the hole, letting a large allocation
                // (e.g. the 1.9 MB display shadow) straddle it.
                let mut lo = s;
                for (b, sz) in holes {
                    if sz > 0 && e > b && s < b + sz {
                        let h0 = b.max(s);
                        let h1 = (b + sz).min(e);
                        if h0 > lo {
                            // Segment before the hole.
                            let seg = h0 - lo;
                            if seg > best_size {
                                best_start = lo;
                                best_size = seg;
                            }
                        }
                        lo = h1;
                    }
                }
                if e > lo {
                    let seg = e - lo;
                    if seg > best_size {
                        best_start = lo;
                        best_size = seg;
                    }
                }
            }
        }
        if best_size > 2 * 1024 * 1024 {
            let s = best_start + 2 * 1024 * 1024;
            let e = best_start + best_size;
            if s < e && e - s >= MIN_BLOCK {
                HEAP_START = s;
                HEAP_END = e;
                HEAP_CAPACITY = e - s;
                // Seed the free list with one block covering the whole heap.
                let block = s as *mut Header;
                (*block).size = e - s;
                (*block).free = 1;
                let links = (*block).links();
                (*links).prev = ptr::null_mut();
                (*links).next = ptr::null_mut();
                FREE_LIST = block;
            }
        }
    }
}

#[allow(dead_code)]
pub fn allocate(size: usize) -> *mut u8 {
    let layout = unsafe { Layout::from_size_align_unchecked(size.max(1), MIN_ALIGN) };
    unsafe { ALLOCATOR.alloc(layout) }
}

#[allow(dead_code)]
pub fn get_map() -> (&'static [MemoryDescriptor], usize) {
    unsafe {
        (
            core::slice::from_raw_parts(&MEMORY_MAP[0], MAP_ENTRIES),
            MAP_KEY,
        )
    }
}

// ---------------------------------------------------------------------------
// Heap statistics + self-test
// ---------------------------------------------------------------------------

/// Snapshot of current heap usage.
#[derive(Debug, Clone, Copy)]
pub struct HeapStats {
    pub heap_start: usize,
    pub heap_end: usize,
    pub heap_size: usize,
    pub free_bytes: usize,
    pub free_blocks: usize,
}

/// Current heap usage — walks the free list (cheap, it stays short).
pub fn heap_stats() -> HeapStats {
    unsafe {
        let mut free_bytes = 0usize;
        let mut free_blocks = 0usize;
        let mut cur = FREE_LIST;
        while !cur.is_null() {
            free_bytes += (*cur).size;
            free_blocks += 1;
            cur = (*(*cur).links()).next;
        }
        HeapStats {
            heap_start: HEAP_START,
            heap_end: HEAP_END,
            heap_size: HEAP_CAPACITY,
            free_bytes,
            free_blocks,
        }
    }
}

/// DEBUG (crash-family): which free block contains `addr`, if any.
/// Returns `(block_base, block_size, offset_within_block)`.
///
/// The OptiPlex heap-corruption dumps land at deterministic offsets
/// inside the shadow's would-be span; this tells us whether the shadow
/// block (or any block spanning that offset) has leaked back into the
/// free list at shell entry or later.
pub fn freelist_contains(addr: usize) -> Option<(usize, usize, usize)> {
    unsafe {
        let mut cur = FREE_LIST;
        while !cur.is_null() {
            let base = cur as usize;
            let size = (*cur).size;
            if addr >= base && addr < base.saturating_add(size) {
                return Some((base, size, addr - base));
            }
            cur = (*(*cur).links()).next;
        }
        None
    }
}

/// DEBUG (crash-family): the free list head block (base, size).
pub fn freelist_head() -> Option<(usize, usize)> {
    unsafe {
        if FREE_LIST.is_null() {
            None
        } else {
            Some((FREE_LIST as usize, (*FREE_LIST).size))
        }
    }
}

/// DEBUG (OptiPlex sim): force the first ExitBootServices call to fail,
/// exercising the capture_memory_map -> init_heap free-list rebuild that
/// wiped the pre-EBS shadow on real hardware.  Returns true when the
/// simulation is armed (one shot).
pub static mut SIMULATE_EBS_FAILURE: bool = false;
pub fn simulate_ebs_failure() -> bool {
    unsafe {
        let r = SIMULATE_EBS_FAILURE;
        SIMULATE_EBS_FAILURE = false;
        r
    }
}

/// DEBUG (crash-family): walk the free list, formatting at most `max`
/// blocks as text lines for the screen/serial.
///
/// Snapshot the walk into a fixed stack array first: printing the lines
/// itself allocates, and an interleaved alloc/free would shift the list
/// under the walk (the shell's tty path allocates per line).
pub fn freelist_lines(max: usize) -> alloc::vec::Vec<alloc::string::String> {
    let mut out = alloc::vec::Vec::new();
    unsafe {
        let mut cur = FREE_LIST;
        let cap = max.min(24);
        let mut snap: [(usize, usize); 24] = [(0, 0); 24];
        let mut seen = 0usize;
        while !cur.is_null() && seen < cap {
            snap[seen] = (cur as usize, (*cur).size);
            seen += 1;
            cur = (*(*cur).links()).next;
        }
        if seen == 0 {
            out.push("fl[] EMPTY".into());
        } else {
            for (base, size) in snap[..seen].iter() {
                out.push(alloc::format!(
                    "fl[{}] {:#x} size={:#x} end={:#x}",
                    out.len(),
                    base,
                    size,
                    base.wrapping_add(*size)
                ));
            }
        }
    }
    out
}

/// Allocator smoke test — exercises alloc/free/coalesce without touching
/// the heap itself (fixed-size stack arrays only).  Returns true on success.
pub fn self_test() -> bool {
    unsafe {
        // Snapshot the heap before we start: earlier boot code (e.g. the
        // vt cell grids) may legitimately hold live allocations.  The test
        // must return the heap to exactly this state, not to "empty".
        let before = heap_stats();

        const N: usize = 128;
        // (ptr, size, align) triplets — fixed array on the stack.
        let mut blocks: [(usize, usize, usize); N] = [(0, 0, 0); N];

        // Allocate a mixed batch: sizes 1..2048, aligns 16..2048.
        let mut seed: u64 = 0x9E37_79B9_7F4A_7C15;
        let mut n = 0usize;
        for i in 0..N {
            let size = ((seed as usize) % 2048) + 1;
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let align = 16usize << (seed as usize % 8);
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let Ok(layout) = Layout::from_size_align(size, align) else {
                return false; // impossible: size <= 2048, align <= 2048
            };
            let p = ALLOCATOR.alloc(layout);
            if p.is_null() {
                break; // heap exhausted — stop
            }
            if !(p as usize).is_multiple_of(align) {
                return false; // alignment violated
            }
            core::ptr::write_bytes(p, i as u8 + 1, size); // scribble
            blocks[n] = (p as usize, size, align);
            n += 1;
        }
        if n < 16 {
            return false; // should easily fit in the heap
        }

        // Free every third block.
        let mut i = 0;
        while i < n {
            let (ptr, size, align) = blocks[i];
            let layout = Layout::from_size_align(size, align).unwrap();
            ALLOCATOR.dealloc(ptr as *mut u8, layout);
            blocks[i] = (0, 0, 0);
            i += 3;
        }

        // Reallocate the freed slots — space must be reused.
        for i in 0..n {
            if blocks[i].0 != 0 {
                continue; // still live
            }
            let size = ((seed as usize) % 2048) + 1;
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let align = 16usize << (seed as usize % 8);
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let Ok(layout) = Layout::from_size_align(size, align) else {
                return false; // impossible: size <= 2048, align <= 2048
            };
            let p = ALLOCATOR.alloc(layout);
            if p.is_null() {
                return false; // failed to reuse freed space
            }
            if !(p as usize).is_multiple_of(align) {
                return false;
            }
            core::ptr::write_bytes(p, 0x5A, size);
            blocks[i] = (p as usize, size, align);
        }

        // Free everything and check the heap coalesces back to one block.
        for (ptr, size, align) in blocks.iter() {
            if *ptr == 0 {
                continue;
            }
            let Ok(layout) = Layout::from_size_align(*size, *align) else {
                return false; // impossible: sizes/aligns came from the loop above
            };
            ALLOCATOR.dealloc(*ptr as *mut u8, layout);
        }
        let stats = heap_stats();
        stats.free_blocks == before.free_blocks && stats.free_bytes == before.free_bytes
    }
}

// ---------------------------------------------------------------------------
// Physical page allocator — bitmap over the usable (conventional) regions.
// ---------------------------------------------------------------------------
// One bit per 4 KB page: 0 = free, 1 = allocated.  The bitmap spans the whole
// [lowest conventional page, highest conventional page] range; pages that are
// NOT conventional inside that span stay marked allocated, so MMIO holes and
// firmware-reserved gaps are never handed out.  The kernel heap region is also
// reserved, keeping the heap and page allocators from double-allocating the
// same RAM.
// ---------------------------------------------------------------------------

pub const PAGE_SIZE: usize = 4096;

/// Bitmap storage (allocated once from the heap at init).
static mut BITMAP: *mut u8 = ptr::null_mut();
/// Number of bytes in the bitmap.
static mut BITMAP_BYTES: usize = 0;
/// Page number that bitmap bit 0 refers to.
static mut BASE_PAGE: usize = 0;
/// Number of pages covered by the bitmap.
static mut SPAN_PAGES: usize = 0;
/// Running count of free pages (kept in sync with the bitmap).
static mut FREE_PAGES: usize = 0;

/// Snapshot of physical page usage.
#[derive(Debug, Clone, Copy)]
pub struct PageStats {
    pub base_page: usize,
    pub span_pages: usize,
    pub span_bytes: usize,
    pub free_pages: usize,
    pub used_pages: usize,
}

/// Current physical page usage.
pub fn page_stats() -> PageStats {
    unsafe {
        PageStats {
            base_page: BASE_PAGE,
            span_pages: SPAN_PAGES,
            span_bytes: SPAN_PAGES * PAGE_SIZE,
            free_pages: FREE_PAGES,
            used_pages: SPAN_PAGES.saturating_sub(FREE_PAGES),
        }
    }
}

/// Build the bitmap from the stashed memory map and reserve the heap.
/// Returns false if there is no conventional memory or the bitmap cannot
/// be allocated.
pub fn init_page_allocator() -> bool {
    unsafe {
        // Lowest and highest conventional page numbers.
        let mut low: usize = usize::MAX;
        let mut high: usize = 0;
        for i in 0..MAP_ENTRIES {
            let d = &MEMORY_MAP[i];
            if d.kind == MEMORY_TYPE_CONVENTIONAL {
                let s = d.physical_start as usize;
                let e = s + d.number_of_pages as usize * PAGE_SIZE;
                if s < low {
                    low = s;
                }
                if e > high {
                    high = e;
                }
            }
        }
        if low == usize::MAX || high <= low {
            return false;
        }

        let base_page = low / PAGE_SIZE;
        let span_pages = high.div_ceil(PAGE_SIZE) - base_page;
        let bytes = span_pages.div_ceil(8);

        // The bitmap itself lives in the heap (one allocation, never freed).
        // `bytes` >= 1 and align 16 is a power of two, so this cannot fail.
        let layout = Layout::from_size_align_unchecked(bytes, 16);
        let bm = ALLOCATOR.alloc(layout);
        if bm.is_null() {
            return false;
        }

        // Start everything allocated, then free only usable pages.
        core::ptr::write_bytes(bm, 0xFF, bytes);

        for i in 0..MAP_ENTRIES {
            let d = &MEMORY_MAP[i];
            if d.kind == MEMORY_TYPE_CONVENTIONAL {
                let s = d.physical_start as usize / PAGE_SIZE;
                let e = (d.physical_start as usize + d.number_of_pages as usize * PAGE_SIZE)
                    .div_ceil(PAGE_SIZE);
                for p in s..e {
                    let idx = p - base_page;
                    *bm.add(idx / 8) &= !(1 << (idx % 8));
                }
            }
        }

        // Reserve the heap so the page allocator never hands out its pages.
        if HEAP_END > HEAP_START {
            let s = HEAP_START / PAGE_SIZE;
            let e = HEAP_END.div_ceil(PAGE_SIZE);
            for p in s..e {
                let idx = p - base_page;
                *bm.add(idx / 8) |= 1 << (idx % 8);
            }
        }

        // Reserve the kernel image so the page allocator never hands out
        // the running kernel's own code/data/bss pages (DMA buffers and
        // task stacks landing on the kernel would corrupt it — a heap
        // allocation over the .bss already broke the FAT ACTIVE_VOL).
        if let Some((ks, ke)) = kernel_image_range()
            && ke > ks
        {
            let s = ks / PAGE_SIZE;
            let e = ke.div_ceil(PAGE_SIZE);
            for p in s..e {
                let idx = p - base_page;
                if idx < span_pages {
                    *bm.add(idx / 8) |= 1 << (idx % 8);
                }
            }
        }

        // Reserve the GOP framebuffer region even if the map lists it as
        // conventional — kernel allocations (stacks, DMA buffers) must
        // never land on the display.
        if let Some((fbase, fsize)) = crate::gop::framebuffer_region()
            && fsize > 0
        {
            let s = fbase / PAGE_SIZE;
            let e = (fbase + fsize).div_ceil(PAGE_SIZE);
            for p in s..e {
                let idx = p - base_page;
                if idx < span_pages {
                    *bm.add(idx / 8) |= 1 << (idx % 8);
                }
            }
        }

        // Never hand out page 0 (the null-pointer trap page).
        if base_page == 0 {
            *bm |= 0x01;
        }

        // Count free pages once; keep the count in sync from here on.
        let mut free = 0usize;
        for i in 0..span_pages {
            if (*bm.add(i / 8) >> (i % 8)) & 1 == 0 {
                free += 1;
            }
        }

        BITMAP = bm;
        BITMAP_BYTES = bytes;
        BASE_PAGE = base_page;
        SPAN_PAGES = span_pages;
        FREE_PAGES = free;
        true
    }
}

/// Allocate `count` contiguous physical pages.  Returns the physical
/// address of the first page, or None when no run fits.
pub fn alloc_pages(count: usize) -> Option<usize> {
    alloc_pages_aligned(count, 1)
}

/// Allocate `count` contiguous physical pages at or above `min_addr`
/// (first-fit upward from there).  Used for controller DMA buffers that
/// must avoid low conventional memory — the first 640 K on some real
/// chipsets is hostile to PCI masters and shows up as a controller
/// "System Error" (SE) halt.  Page-aligned like `alloc_pages`.
pub fn alloc_pages_above(count: usize, min_addr: usize) -> Option<usize> {
    alloc_pages_range(count, min_addr, usize::MAX)
}

/// Allocate `count` contiguous physical pages inside `[min_addr, max_addr)`.
/// The upper bound matters for 32-bit DMA engines (EHCI): a buffer above
/// 4 GiB is unreachable even though the kernel can address it.
pub fn alloc_pages_range(count: usize, min_addr: usize, max_addr: usize) -> Option<usize> {
    if count == 0 {
        return None;
    }
    unsafe {
        if BITMAP.is_null() {
            return None;
        }
        let span = SPAN_PAGES;
        let min_page = min_addr.div_ceil(PAGE_SIZE).max(BASE_PAGE);
        let max_page = (max_addr / PAGE_SIZE).min(BASE_PAGE + span);
        if min_page >= max_page {
            return None;
        }
        let start_idx = min_page - BASE_PAGE;
        let end_idx = max_page - BASE_PAGE;
        let mut run_start = 0usize;
        let mut run_len = 0usize;
        for idx in start_idx..end_idx {
            let free = (*BITMAP.add(idx / 8) >> (idx % 8)) & 1 == 0;
            if free {
                if run_len == 0 {
                    run_start = idx;
                }
                run_len += 1;
                if run_len >= count {
                    let p = (BASE_PAGE + run_start) * PAGE_SIZE;
                    for k in run_start..run_start + count {
                        *BITMAP.add(k / 8) |= 1 << (k % 8);
                    }
                    FREE_PAGES -= count;
                    return Some(p);
                }
            } else {
                run_len = 0;
            }
        }
        None
    }
}

/// Allocate `count` contiguous physical pages whose start address is
/// aligned to `align_pages` (a power of two, in page units).
pub fn alloc_pages_aligned(count: usize, align_pages: usize) -> Option<usize> {
    if count == 0 || !align_pages.is_power_of_two() {
        return None;
    }
    unsafe {
        if BITMAP.is_null() || FREE_PAGES < count {
            return None;
        }
        let span = SPAN_PAGES;
        let mask = align_pages - 1;
        let mut run_start = 0usize;
        let mut run_len = 0usize;
        for idx in 0..span {
            let free = (*BITMAP.add(idx / 8) >> (idx % 8)) & 1 == 0;
            if free {
                if run_len == 0 {
                    run_start = idx;
                }
                run_len += 1;
                // Aligned page number within (or just past) this run.
                let p = BASE_PAGE + run_start;
                let p_align = (p + mask) & !mask;
                if p_align == 0 {
                    // Candidate starts at page 0 — restart the run here.
                    run_start = idx;
                    run_len = 1;
                    continue;
                }
                let idx_align = p_align - BASE_PAGE;
                if idx_align + count <= run_start + run_len {
                    // Mark the run allocated.
                    for k in idx_align..idx_align + count {
                        *BITMAP.add(k / 8) |= 1 << (k % 8);
                    }
                    FREE_PAGES -= count;
                    return Some(p_align * PAGE_SIZE);
                }
            } else {
                run_len = 0;
            }
        }
        None
    }
}

/// Free `count` physical pages starting at `addr`.  The range must be
/// page-aligned and previously allocated by `alloc_pages`.
pub fn free_pages(addr: usize, count: usize) {
    unsafe {
        if BITMAP.is_null() || count == 0 {
            return;
        }
        let base = BASE_PAGE;
        let span = SPAN_PAGES;
        let first = addr / PAGE_SIZE;
        let in_range =
            addr.is_multiple_of(PAGE_SIZE) && first >= base && first + count <= base + span;
        if !in_range {
            debug_assert!(false, "free_pages: bad range {:#x} len {}", addr, count);
            return;
        }
        for p in first..first + count {
            let idx = p - base;
            debug_assert!(
                (*BITMAP.add(idx / 8) >> (idx % 8)) & 1 != 0,
                "free_pages: double free at {:#x}",
                p * PAGE_SIZE
            );
            *BITMAP.add(idx / 8) &= !(1 << (idx % 8));
        }
        FREE_PAGES += count;
    }
}

/// Smoke test for the page allocator: distinct single pages, a contiguous
/// run, scribble/readback, and a clean return to the baseline free count.
pub fn page_self_test() -> bool {
    unsafe {
        let before = page_stats();
        if before.free_pages < 1024 {
            return false; // not enough free pages to play with
        }

        // 64 single pages — all distinct, page-aligned, non-zero.
        let mut pages: [usize; 64] = [0; 64];
        for i in 0..64 {
            let Some(addr) = alloc_pages(1) else {
                return false;
            };
            if addr % PAGE_SIZE != 0 || addr == 0 {
                return false;
            }
            pages[i] = addr;
        }
        for i in 0..64 {
            for j in 0..i {
                if pages[i] == pages[j] {
                    return false; // overlap
                }
            }
        }

        // A contiguous run of 16.
        let Some(run) = alloc_pages(16) else {
            return false;
        };
        if run % PAGE_SIZE != 0 || run == 0 {
            return false;
        }

        // Free count must drop by exactly 64 + 16.
        if page_stats().free_pages != before.free_pages - 64 - 16 {
            return false;
        }

        // Scribble through the run and read it back (it's real RAM).
        for k in 0..16 {
            core::ptr::write_volatile((run + k * PAGE_SIZE) as *mut u8, 0xAB);
        }
        for k in 0..16 {
            if core::ptr::read_volatile((run + k * PAGE_SIZE) as *const u8) != 0xAB {
                return false;
            }
        }

        // Free everything and check we're back to baseline.
        for &p in pages.iter() {
            free_pages(p, 1);
        }
        free_pages(run, 16);

        let after = page_stats();
        after.free_pages == before.free_pages
    }
}