//! X64 Paging
//!
//! This module provides an in direction to the external paging/mtrr crates.
//!
//! ## License
//!
//! Copyright (c) Microsoft Corporation.
//!
//! SPDX-License-Identifier: Apache-2.0
//!
use crate::paging::{CacheAttributeValue, PatinaPageTable};
use alloc::boxed::Box;
use patina::error::EfiError;
use patina_mtrr::{Mtrr, create_mtrr_lib, error::MtrrError, structs::MtrrMemoryCacheType};
use patina_paging::{
    MemoryAttributes, PageTable, PagingType, PtError, page_allocator::PageAllocator, x64::X64PageTable,
};
use r_efi::efi;

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

fn efierror_to_pterror(efi_error: EfiError) -> PtError {
    match efi_error {
        EfiError::InvalidParameter => PtError::InvalidParameter,
        EfiError::OutOfResources => PtError::OutOfResources,
        EfiError::NotFound => PtError::NoMapping,
        _ => PtError::InvalidParameter, // Default case for unsupported error codes
    }
}

/// The x86_64 paging implementation.
impl<P, M> PatinaPageTable for EfiCpuPagingX64<P, M>
where
    P: PageTable,
    M: Mtrr,
{
    // Paging related APIs
    fn map_memory_region(&mut self, address: u64, size: u64, attributes: MemoryAttributes) -> Result<(), PtError> {
        let cache_attributes = attributes & MemoryAttributes::CacheAttributesMask;
        let memory_attributes = attributes & MemoryAttributes::AccessAttributesMask;

        if attributes != (cache_attributes | memory_attributes) {
            log::error!("Invalid cache attribute: {:#x}", attributes);
            return Err(PtError::InvalidParameter);
        }

        match apply_caching_attributes(address, size, cache_attributes, &mut self.mtrr) {
            Ok(_) => self.paging.map_memory_region(address, size, attributes & MemoryAttributes::AccessAttributesMask),
            Err(status) => Err(efierror_to_pterror(status)),
        }
    }

    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)> {
        // start by getting the caching attributes as we need to return those even if the page is unmapped in the
        // page table
        let cache_attr = match self.mtrr.get_memory_attribute(address) {
            MtrrMemoryCacheType::Uncacheable => CacheAttributeValue::Valid(MemoryAttributes::Uncacheable),
            MtrrMemoryCacheType::WriteCombining => CacheAttributeValue::Valid(MemoryAttributes::WriteCombining),
            MtrrMemoryCacheType::WriteThrough => CacheAttributeValue::Valid(MemoryAttributes::WriteThrough),
            MtrrMemoryCacheType::WriteProtected => CacheAttributeValue::Valid(MemoryAttributes::WriteProtect),
            MtrrMemoryCacheType::WriteBack => CacheAttributeValue::Valid(MemoryAttributes::Writeback),
            _ => CacheAttributeValue::Unmapped,
        };

        match self.paging.query_memory_region(address, size) {
            Ok(attr) => match cache_attr {
                CacheAttributeValue::Valid(cache_attr_val) => Ok(attr | cache_attr_val),
                _ => {
                    debug_assert!(false, "Cache attributes should be valid for mapped region");
                    Ok(attr)
                }
            },
            Err(err) => Err((err, cache_attr)),
        }
    }

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

fn apply_caching_attributes<M: Mtrr>(
    base_address: u64,
    length: u64,
    cache_attributes: MemoryAttributes,
    mtrr: &mut M,
) -> Result<(), EfiError> {
    if cache_attributes.bits() != 0 {
        if !mtrr.is_supported() {
            return Err(EfiError::Unsupported);
        }

        let cache_type = match cache_attributes {
            MemoryAttributes::Uncacheable => MtrrMemoryCacheType::Uncacheable,
            MemoryAttributes::WriteCombining => MtrrMemoryCacheType::WriteCombining,
            MemoryAttributes::WriteThrough => MtrrMemoryCacheType::WriteThrough,
            MemoryAttributes::WriteProtect => MtrrMemoryCacheType::WriteProtected,
            MemoryAttributes::Writeback => MtrrMemoryCacheType::WriteBack,
            _ => return Err(EfiError::Unsupported),
        };

        let curr_attribute = mtrr.get_memory_attribute(base_address);
        if curr_attribute != cache_type {
            // cache attributes are not already set
            match mtrr.set_memory_attribute(base_address, length, cache_type) {
                Ok(_) => {
                    // now we need to program the APs with the update, if they are up
                    return Ok(());
                }
                Err(err) => return Err(mtrr_err_to_efi_status(err)),
            }
        }
    }

    Ok(())
}

/// Create an x86_64 paging instance under the general PatinaPageTable trait.
pub fn create_cpu_x64_paging<A: PageAllocator + 'static>(
    page_allocator: A,
) -> Result<Box<dyn PatinaPageTable>, efi::Status> {
    Ok(Box::new(EfiCpuPagingX64 {
        paging: X64PageTable::new(page_allocator, PagingType::Paging4Level).unwrap(),
        mtrr: create_mtrr_lib(0),
    }))
}

fn mtrr_err_to_efi_status(err: MtrrError) -> EfiError {
    match err {
        MtrrError::MtrrNotSupported => EfiError::Unsupported,
        MtrrError::VariableRangeMtrrExhausted => EfiError::OutOfResources,
        MtrrError::FixedRangeMtrrBaseAddressNotAligned => EfiError::InvalidParameter,
        MtrrError::FixedRangeMtrrLengthNotAligned => EfiError::InvalidParameter,
        MtrrError::InvalidParameter => EfiError::InvalidParameter,
        MtrrError::BufferTooSmall => EfiError::BufferTooSmall,
        MtrrError::OutOfResources => EfiError::OutOfResources,
        MtrrError::AlreadyStarted => EfiError::AlreadyStarted,
    }
}

#[cfg(test)]
#[coverage(off)]
mod tests {
    use super::*;
    use mockall::mock;
    use patina_mtrr::{
        error::MtrrResult,
        structs::{MtrrMemoryRange, MtrrSettings},
    };

    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>;
        }
    }

    mock! {
        Mtrr {}
        impl Mtrr for Mtrr {
            fn is_supported(&self) -> bool;
            fn get_all_mtrrs(&self) -> MtrrResult<MtrrSettings>;
            fn set_all_mtrrs(&mut self, mtrr_setting: &MtrrSettings);
            fn get_memory_attribute(&self, address: u64) -> MtrrMemoryCacheType;
            fn set_memory_attribute(
                &mut self,
                base_address: u64,
                length: u64,
                attribute: MtrrMemoryCacheType,
            ) -> MtrrResult<()>;
            fn set_memory_attributes(&mut self, ranges: &[MtrrMemoryRange]) -> MtrrResult<()>;
            fn get_memory_ranges(&self) -> MtrrResult<Vec<MtrrMemoryRange>>;

            fn debug_print_all_mtrrs(&self);
        }
    }

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

        mock_page_table.expect_map_memory_region().returning(|_, _, _| Ok(()));
        mock_mtrr.expect_is_supported().return_const(true);
        mock_mtrr.expect_get_memory_attribute().return_const(MtrrMemoryCacheType::Uncacheable);
        mock_mtrr.expect_set_memory_attribute().returning(|_, _, _| Ok(()));

        let mut paging = EfiCpuPagingX64 { paging: mock_page_table, mtrr: mock_mtrr };

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

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

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

        let mut paging = EfiCpuPagingX64 { paging: mock_page_table, mtrr: mock_mtrr };

        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();
        let mut mock_mtrr = MockMtrr::new();

        mock_page_table.expect_map_memory_region().returning(|_, _, _| Ok(()));
        mock_mtrr.expect_is_supported().return_const(true);
        mock_mtrr.expect_get_memory_attribute().return_const(MtrrMemoryCacheType::Uncacheable);
        mock_mtrr.expect_set_memory_attribute().returning(|_, _, _| Ok(()));

        let mut paging = EfiCpuPagingX64 { paging: mock_page_table, mtrr: mock_mtrr };

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

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

        mock_page_table.expect_query_memory_region().returning(|_, _| Ok(MemoryAttributes::Writeback));
        mock_mtrr.expect_get_memory_attribute().return_const(MtrrMemoryCacheType::Uncacheable);

        let paging = EfiCpuPagingX64 { paging: mock_page_table, mtrr: mock_mtrr };

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