use crate::uefi::*;
use core::alloc::{GlobalAlloc, Layout};
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];

// ---------------------------------------------------------------------------
// Bump allocator — no deallocation, just pointer advancement.
// Used as Rust's #[global_allocator] once the heap region is set up.
// ---------------------------------------------------------------------------
struct BumpAllocator;

unsafe impl Sync for BumpAllocator {}

unsafe impl GlobalAlloc for BumpAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let align = layout.align().max(16);
        let size = layout.size();
        unsafe {
            let a = (HEAP_CURSOR + align - 1) & !(align - 1);
            let e = a + size;
            if e > HEAP_END || e < a {
                return ptr::null_mut();
            }
            HEAP_CURSOR = e;
            a as *mut u8
        }
    }

    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
        // Bump allocator — cannot free individual allocations.
        // Memory is reclaimed on reboot.  Good enough for now.
        unsafe {
            core::ptr::read_volatile(&_ptr);
        }
        let _ = _layout;
    }
}

#[global_allocator]
static ALLOCATOR: BumpAllocator = BumpAllocator;

/// Heap bump allocator bounds (`static mut` accessed only inside
/// `GlobalAlloc` impl, which is already `unsafe`).
static mut HEAP_START: usize = 0;
static mut HEAP_END: usize = 0;
static mut HEAP_CURSOR: usize = 0;

/// 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> {
    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 + desc_size - 1) / desc_size * desc_size;

    // If our static buffer is too small, fall back to pool (rare)
    if buf_size > unsafe { 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.

    init_bump_allocator();
    Ok(map_key)
}

/// Retry ExitBootServices with a fresh memory map.
pub fn refresh_and_exit(bs: &BootServices, image_handle: Handle) -> Result<usize, Status> {
    let key = capture_memory_map(bs)?;
    let status = unsafe { (bs.exit_boot_services)(image_handle, key) };
    if status == EFI_SUCCESS {
        Ok(key)
    } else {
        Err(status)
    }
}

fn init_bump_allocator() {
    unsafe {
        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 sz = d.number_of_pages as usize * 4096;
                if sz > best_size {
                    best_start = d.physical_start as usize;
                    best_size = sz;
                }
            }
        }
        if best_size > 2 * 1024 * 1024 {
            let s = best_start + 2 * 1024 * 1024;
            let e = best_start + best_size;
            if s < e {
                HEAP_START = s;
                HEAP_END = e;
                HEAP_CURSOR = s;
            }
        }
    }
}

#[allow(dead_code)]
pub fn allocate(size: usize) -> *mut u8 {
    unsafe {
        let a = (HEAP_CURSOR + 15) & !15;
        let e = a + size;
        if e > HEAP_END || e < a {
            return ptr::null_mut();
        }
        HEAP_CURSOR = e;
        a as *mut u8
    }
}

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