use x86_64::PhysAddr;
use x86_64::VirtAddr;
use x86_64::structures::paging::{
    Mapper, OffsetPageTable, Page, PageTableFlags, PhysFrame, Size4KiB,
};

use super::phys::BitmapFrameAllocator;

pub fn map_mmio(
    mapper: &mut OffsetPageTable,
    allocator: &mut BitmapFrameAllocator,
    phys: u64,
    hhdm_offset: u64,
) -> Result<(), crate::error::KernelError> {
    let virt = VirtAddr::new(phys + hhdm_offset);
    let page: Page<Size4KiB> = Page::containing_address(virt);
    let frame = PhysFrame::containing_address(PhysAddr::new(phys));
    let flags = PageTableFlags::PRESENT
        | PageTableFlags::WRITABLE
        | PageTableFlags::NO_CACHE
        | PageTableFlags::WRITE_THROUGH
        | PageTableFlags::NO_EXECUTE;

    match unsafe { mapper.map_to(page, frame, flags, allocator) } {
        Ok(flush) => {
            flush.flush();
            Ok(())
        }
        Err(x86_64::structures::paging::mapper::MapToError::PageAlreadyMapped(existing)) => {
            if existing == frame {
                Ok(())
            } else {
                Err(crate::error::KernelError::InvalidAddress)
            }
        }
        Err(_) => Err(crate::error::KernelError::ResourceExhausted),
    }
}

pub fn test_mapping(mapper: &mut OffsetPageTable, allocator: &mut BitmapFrameAllocator) {
    let test_addr = VirtAddr::new(0xFFFF_FFFF_F000_0000);
    let test_page: Page<Size4KiB> = Page::containing_address(test_addr);

    let frame = allocator.allocate().expect("Failed to allocate test frame");

    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
    let phys_frame = frame.inner();

    unsafe {
        mapper
            .map_to(test_page, phys_frame, flags, allocator)
            .expect("Failed to map test page")
            .flush();
    }

    let ptr: *mut u64 = test_addr.as_mut_ptr();
    unsafe { ptr.write_volatile(0xDEAD_BEEF_CAFE_BABE) };
    let readback = unsafe { ptr.read_volatile() };
    assert_eq!(
        readback, 0xDEAD_BEEF_CAFE_BABE,
        "Virtual memory test failed"
    );

    let (returned_frame, flush) = mapper.unmap(test_page).expect("Failed to unmap test page");
    flush.flush();
    allocator.deallocate_frame(returned_frame);

    crate::kprintln!("  Virtual memory mapping test passed");
}