use bitflags::bitflags;

/// Describes a single page in the page table.
/// Mirrors the C++ `PageEntry` union from memory.h:82-99.
#[derive(Clone, Copy, Default)]
pub struct PageEntry(u64);

impl PageEntry {
    /// Reconstruct a [`PageEntry`] from its packed `u64` representation.
    /// Used by [`crate::GuestMemory::is_mapped`] and `page_entry` after an
    /// atomic load from the page table.
    pub fn from_raw(raw: u64) -> Self {
        Self(raw)
    }

    /// The packed `u64` representation, ready to atomically Release-store
    /// into the page table.
    pub fn raw(&self) -> u64 {
        self.0
    }

    /// Base address of the allocated region in 4K pages (20 bits).
    pub fn base_address(&self) -> u32 {
        (self.0 & 0xFFFFF) as u32
    }

    pub fn set_base_address(&mut self, val: u32) {
        self.0 = (self.0 & !0xFFFFF) | (val as u64 & 0xFFFFF);
    }

    /// Total number of pages in the allocated region (20 bits).
    pub fn region_page_count(&self) -> u32 {
        ((self.0 >> 20) & 0xFFFFF) as u32
    }

    pub fn set_region_page_count(&mut self, val: u32) {
        self.0 = (self.0 & !(0xFFFFF << 20)) | ((val as u64 & 0xFFFFF) << 20);
    }

    /// Protection bits specified during region allocation (4 bits).
    pub fn allocation_protect(&self) -> MemoryProtect {
        MemoryProtect::from_bits_truncate(((self.0 >> 40) & 0xF) as u32)
    }

    pub fn set_allocation_protect(&mut self, val: MemoryProtect) {
        self.0 = (self.0 & !(0xF << 40)) | ((val.bits() as u64 & 0xF) << 40);
    }

    /// Current protection bits (4 bits).
    pub fn current_protect(&self) -> MemoryProtect {
        MemoryProtect::from_bits_truncate(((self.0 >> 44) & 0xF) as u32)
    }

    pub fn set_current_protect(&mut self, val: MemoryProtect) {
        self.0 = (self.0 & !(0xF << 44)) | ((val.bits() as u64 & 0xF) << 44);
    }

    /// Allocation state (2 bits).
    pub fn state(&self) -> AllocationState {
        AllocationState::from_bits_truncate(((self.0 >> 48) & 0x3) as u32)
    }

    pub fn set_state(&mut self, val: AllocationState) {
        self.0 = (self.0 & !(0x3 << 48)) | ((val.bits() as u64 & 0x3) << 48);
    }

    pub fn is_committed(&self) -> bool {
        self.state().contains(AllocationState::COMMIT)
    }

    pub fn is_reserved(&self) -> bool {
        self.state().contains(AllocationState::RESERVE)
    }

    pub fn is_free(&self) -> bool {
        self.state().is_empty()
    }
}

bitflags! {
    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
    pub struct MemoryProtect: u32 {
        const READ = 1 << 0;
        const WRITE = 1 << 1;
        const NO_CACHE = 1 << 2;
        const WRITE_COMBINE = 1 << 3;
    }
}

bitflags! {
    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
    pub struct AllocationState: u32 {
        const RESERVE = 1 << 0;
        const COMMIT = 1 << 1;
    }
}

impl std::fmt::Debug for PageEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PageEntry")
            .field("base_address", &format_args!("{:#x}", self.base_address()))
            .field("region_page_count", &self.region_page_count())
            .field("allocation_protect", &self.allocation_protect())
            .field("current_protect", &self.current_protect())
            .field("state", &self.state())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_page_entry_bitfields() {
        let mut entry = PageEntry::default();
        assert!(entry.is_free());

        entry.set_base_address(0x100);
        entry.set_region_page_count(0x10);
        entry.set_allocation_protect(MemoryProtect::READ | MemoryProtect::WRITE);
        entry.set_current_protect(MemoryProtect::READ);
        entry.set_state(AllocationState::RESERVE | AllocationState::COMMIT);

        assert_eq!(entry.base_address(), 0x100);
        assert_eq!(entry.region_page_count(), 0x10);
        assert_eq!(
            entry.allocation_protect(),
            MemoryProtect::READ | MemoryProtect::WRITE
        );
        assert_eq!(entry.current_protect(), MemoryProtect::READ);
        assert!(entry.is_committed());
        assert!(entry.is_reserved());
    }
}