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

use crate::mem::phys::BitmapFrameAllocator;

crate::kernel_test!(
    fn map_write_read_unmap() {
        let hhdm = crate::mem::addr::hhdm_offset();
        let mut mapper = unsafe { crate::arch::paging::init(hhdm) };
        let mut alloc = BitmapFrameAllocator;

        let test_va = VirtAddr::new(0xFFFF_FFFF_E000_0000);
        let page: Page<Size4KiB> = Page::containing_address(test_va);
        let frame = alloc.allocate().expect("alloc frame for map test");
        let phys_frame = frame.inner();
        let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

        unsafe {
            mapper
                .map_to(page, phys_frame, flags, &mut alloc)
                .expect("map_to failed")
                .flush();
        }

        let ptr: *mut u64 = test_va.as_mut_ptr();
        let pattern: u64 = 0x1234_5678_9ABC_DEF0;
        unsafe { ptr.write_volatile(pattern) };
        let readback = unsafe { ptr.read_volatile() };
        assert_eq!(readback, pattern, "paging write/read mismatch");

        let (returned, flush) = mapper.unmap(page).expect("unmap failed");
        flush.flush();
        alloc.deallocate_frame(returned);
    }
);

crate::kernel_test!(
    fn dual_mapping_shares_physical() {
        let hhdm = crate::mem::addr::hhdm_offset();
        let mut mapper = unsafe { crate::arch::paging::init(hhdm) };
        let mut alloc = BitmapFrameAllocator;

        let va_a = VirtAddr::new(0xFFFF_FFFF_D000_0000);
        let va_b = VirtAddr::new(0xFFFF_FFFF_D000_1000);
        let page_a: Page<Size4KiB> = Page::containing_address(va_a);
        let page_b: Page<Size4KiB> = Page::containing_address(va_b);

        let frame = alloc.allocate().expect("alloc for dual-map");
        let phys = frame.inner();
        let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

        unsafe {
            mapper
                .map_to(page_a, phys, flags, &mut alloc)
                .expect("map_to A")
                .flush();
            mapper
                .map_to(page_b, phys, flags, &mut alloc)
                .expect("map_to B")
                .flush();
        }

        let ptr_a: *mut u64 = va_a.as_mut_ptr();
        let ptr_b: *const u64 = va_b.as_ptr();
        let pattern: u64 = 0xCAFE_BABE_F00D_FACE;
        unsafe { ptr_a.write_volatile(pattern) };
        let read_b = unsafe { ptr_b.read_volatile() };
        assert_eq!(read_b, pattern, "dual-mapped read mismatch");

        let (ret_a, flush_a) = mapper.unmap(page_a).expect("unmap A");
        flush_a.flush();
        let (_, flush_b) = mapper.unmap(page_b).expect("unmap B");
        flush_b.flush();
        alloc.deallocate_frame(ret_a);
    }
);