//! AArch64 Paging
//!
//! This module provides an in direction to the external paging crate.
//!
//! ## License
//!
//! Copyright (c) Microsoft Corporation.
//!
//! SPDX-License-Identifier: Apache-2.0
//!
use patina_paging::{MemoryAttributes, PageTable, PagingType, PtError, aarch64::AArch64PageTable};

use crate::{
    cpu::aarch64::flush_data_cache_range,
    paging::{CacheAttributeValue, PatinaPageTable},
};
use patina::pi::protocols::cpu_arch::CpuFlushType;
use patina_paging::page_allocator::PageAllocator;
use r_efi::efi;

/// Memory attributes that indicate cached mappings (writeback or write-through).
const CACHED_ATTRS: MemoryAttributes = MemoryAttributes::Writeback.union(MemoryAttributes::WriteThrough);

/// Memory attributes that indicate uncached mappings (uncached or write-combining).
const UNCACHED_ATTRS: MemoryAttributes = MemoryAttributes::Uncached.union(MemoryAttributes::WriteCombining);

/// The aarch64 paging implementation. It acts as a bridge between the EFI CPU
/// Architecture Protocol and the aarch64 paging implementation.
#[derive(Debug)]
pub struct EfiCpuPagingAArch64<P>
where
    P: PageTable,
{
    paging: P,
}

/// The aarch64 paging implementation.
impl<P> PatinaPageTable for EfiCpuPagingAArch64<P>
where
    P: PageTable,
{
    fn map_memory_region(&mut self, address: u64, size: u64, attributes: MemoryAttributes) -> Result<(), PtError> {
        self.paging.map_memory_region(address, size, attributes)
    }

    fn unmap_memory_region(&mut self, address: u64, size: u64) -> Result<(), PtError> {
        self.paging.unmap_memory_region(address, size)
    }

    fn install_page_table(&mut self) -> Result<(), PtError> {
        self.paging.install_page_table()
    }

    fn query_memory_region(&self, address: u64, size: u64) -> Result<MemoryAttributes, (PtError, CacheAttributeValue)> {
        // in AARCH64, the caching attributes are managed in the page table and so we will never return just caching
        // attributes
        self.paging.query_memory_region(address, size).map_err(|e| (e, CacheAttributeValue::NotSupported))
    }

    fn dump_page_tables(&self, address: u64, size: u64) -> Result<(), PtError> {
        self.paging.dump_page_tables(address, size)
    }

    fn handle_cacheability_change(
        &self,
        address: u64,
        size: u64,
        old_cache_attributes: MemoryAttributes,
        new_cache_attributes: MemoryAttributes,
    ) {
        // If the region was previously cached (WB/WT) and is now uncached (UC/WC),
        // clean and invalidate the data cache for the range so that any dirty lines
        // are written back to memory before subsequent accesses bypass the cache.
        if old_cache_attributes.intersects(CACHED_ATTRS) && new_cache_attributes.intersects(UNCACHED_ATTRS) {
            flush_data_cache_range(address, size, CpuFlushType::EfiCpuFlushTypeWriteBackInvalidate);
        }
    }
}

/// Create an AArch64 paging instance under the general PatinaPageTable trait.
#[coverage(off)]
pub fn create_cpu_aarch64_paging<A: PageAllocator + 'static>(
    page_allocator: A,
) -> Result<impl PatinaPageTable, efi::Status> {
    Ok(EfiCpuPagingAArch64 { paging: AArch64PageTable::new(page_allocator, PagingType::Paging4Level).unwrap() })
}

/// Open the active AArch64 page table wrapped in the PatinaPageTable trait.
///
/// ## Safety
/// The caller must ensure no other entity is concurrently modifying the page tables.
#[coverage(off)]
pub unsafe fn open_active_cpu_aarch64_paging<A: PageAllocator + 'static>(
    page_allocator: A,
) -> Result<impl PatinaPageTable, PtError> {
    // SAFETY: Caller ensures no concurrent page table modifications.
    let page_table = unsafe { AArch64PageTable::open_active(page_allocator)? };
    Ok(EfiCpuPagingAArch64 { paging: page_table })
}

#[cfg(test)]
#[coverage(off)]
mod tests {

    use super::*;
    use mockall::mock;

    mock! {
        PageAllocator {}
        impl PageAllocator for PageAllocator {
            fn allocate_page(&mut self, align: u64, size: u64, is_root: bool) -> Result<u64, PtError>;
        }
    }

    mock! {
        PageTable {}
        impl PageTable for PageTable {
            fn map_memory_region(&mut self, address: u64, size: u64, attributes: MemoryAttributes) -> Result<(), PtError>;
            fn unmap_memory_region(&mut self, address: u64, size: u64) -> Result<(), PtError>;
            fn install_page_table(&mut self) -> Result<(), PtError>;
            fn query_memory_region(&self, address: u64, size: u64) -> Result<MemoryAttributes, PtError>;
            fn dump_page_tables(&self, address: u64, size: u64) -> Result<(), PtError>;
        }
    }

    #[test]
    fn test_map_memory_region() {
        let mut mock_page_table = MockPageTable::new();

        mock_page_table.expect_map_memory_region().returning(|_, _, _| Ok(()));

        let mut paging = EfiCpuPagingAArch64 { paging: mock_page_table };

        let result = paging.map_memory_region(0x1000, 0x1000, MemoryAttributes::Uncached);
        assert!(result.is_ok());
    }

    #[test]
    fn test_unmap_memory_region() {
        let mut mock_page_table = MockPageTable::new();

        mock_page_table.expect_unmap_memory_region().returning(|_, _| Ok(()));

        let mut paging = EfiCpuPagingAArch64 { paging: mock_page_table };

        let result = paging.unmap_memory_region(0x1000, 0x1000);
        assert!(result.is_ok());
    }

    #[test]
    fn test_remap_memory_region() {
        let mut mock_page_table = MockPageTable::new();

        mock_page_table.expect_map_memory_region().returning(|_, _, _| Ok(()));

        let mut paging = EfiCpuPagingAArch64 { paging: mock_page_table };

        let result = paging.map_memory_region(0x1000, 0x1000, MemoryAttributes::Uncached);
        assert!(result.is_ok());
    }

    #[test]
    fn test_query_memory_region() {
        let mut mock_page_table = MockPageTable::new();

        mock_page_table
            .expect_query_memory_region()
            .returning(|_, _| Ok(MemoryAttributes::Writeback | MemoryAttributes::Uncached));

        let paging = EfiCpuPagingAArch64 { paging: mock_page_table };

        let result = paging.query_memory_region(0x1000, 0x1000);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), MemoryAttributes::Writeback | MemoryAttributes::Uncached);
    }

    #[test]
    fn test_handle_cacheability_change() {
        let paging = EfiCpuPagingAArch64 { paging: MockPageTable::new() };
        paging.handle_cacheability_change(0x1000, 0x1000, MemoryAttributes::Writeback, MemoryAttributes::Uncached);
    }
}
