//! DXE Core CPU Architectural Protocol
//!
//! ## License
//!
//! Copyright (c) Microsoft Corporation.
//!
//! SPDX-License-Identifier: Apache-2.0
//!
#![allow(unused)]
/// Architecture independent public C EFI CPU Architectural Protocol definition.
use crate::{dxe_services, protocols::PROTOCOL_DB};
use alloc::boxed::Box;
use core::ffi::c_void;
use patina::{
    boot_services::{BootServices, StandardBootServices},
    component::{
        Storage, component,
        service::{IntoService, Service},
    },
    error::{EfiError, Result},
    uefi_protocol::ProtocolInterface,
};
use patina_internal_cpu::{
    cpu::{Cpu, EfiCpu},
    interrupts::{self, ExceptionType, HandlerType, InterruptManager, Interrupts},
};
use r_efi::efi;

use patina::pi::protocols::cpu_arch::{CpuFlushType, CpuInitType, InterruptHandler, PROTOCOL_GUID, Protocol};

#[derive(IntoService)]
#[service(dyn Cpu)]
pub(crate) struct DxeCpu(pub(crate) EfiCpu);

impl Cpu for DxeCpu {
    fn flush_data_cache(&self, start: efi::PhysicalAddress, length: u64, flush_type: CpuFlushType) -> Result<()> {
        self.0.flush_data_cache(start, length, flush_type)
    }

    fn init(&self, init_type: CpuInitType) -> Result<()> {
        self.0.init(init_type)
    }

    fn get_timer_value(&self, timer_index: u32) -> Result<(u64, u64)> {
        self.0.get_timer_value(timer_index)
    }

    fn cache_writeback_granule(&self) -> u32 {
        self.0.cache_writeback_granule()
    }
}

#[derive(IntoService)]
#[service(dyn InterruptManager)]
pub(crate) struct DxeInterruptManager(pub(crate) Interrupts);

impl InterruptManager for DxeInterruptManager {
    fn register_exception_handler(&self, exception_type: ExceptionType, handler: HandlerType) -> Result<()> {
        self.0.register_exception_handler(exception_type, handler)
    }

    fn unregister_exception_handler(&self, exception_type: ExceptionType) -> Result<()> {
        self.0.unregister_exception_handler(exception_type)
    }
}

#[repr(C)]
struct EfiCpuArchProtocolImpl {
    protocol: Protocol,

    // Crate accessible fields
    pub(crate) cpu: Service<dyn Cpu>,
    pub(crate) interrupt_manager: Service<dyn InterruptManager>,
}

// SAFETY: EfiCpuArchProtocolImpl provides a valid protocol structure with stable GUID.
unsafe impl ProtocolInterface for EfiCpuArchProtocolImpl {
    const PROTOCOL_GUID: patina::BinaryGuid = PROTOCOL_GUID;
}

// Helper to convert a raw protocol pointer to a reference. Returns `None` when the caller passes a
// null pointer so the caller can determine the appropriate action to take.
fn get_impl_ref<'a>(this: *const Protocol) -> Option<&'a EfiCpuArchProtocolImpl> {
    if this.is_null() {
        return None;
    }

    // SAFETY: `this` is non-null and points to an EfiCpuArchProtocolImpl instance installed by
    //         Patina via `Box::leak`, so it is properly aligned and valid for the protocol's
    //         lifetime.
    Some(unsafe { &*(this as *const EfiCpuArchProtocolImpl) })
}

fn get_impl_ref_mut<'a>(this: *mut Protocol) -> Option<&'a mut EfiCpuArchProtocolImpl> {
    if this.is_null() {
        return None;
    }

    // SAFETY: `this` is non-null and points to an EfiCpuArchProtocolImpl instance installed by
    //         Patina via `Box::leak`, so it is properly aligned and valid for the protocol's
    //         lifetime.
    Some(unsafe { &mut *(this as *mut EfiCpuArchProtocolImpl) })
}

// EfiCpuArchProtocolImpl function pointers implementations.

extern "efiapi" fn flush_data_cache(
    this: *const Protocol,
    start: efi::PhysicalAddress,
    length: u64,
    flush_type: CpuFlushType,
) -> efi::Status {
    let Some(impl_ref) = get_impl_ref(this) else {
        return efi::Status::INVALID_PARAMETER;
    };

    let result = impl_ref.cpu.flush_data_cache(start, length, flush_type);

    result.map(|_| efi::Status::SUCCESS).unwrap_or_else(|err| err.into())
}

extern "efiapi" fn enable_interrupt(this: *const Protocol) -> efi::Status {
    interrupts::enable_interrupts();

    efi::Status::SUCCESS
}

extern "efiapi" fn disable_interrupt(this: *const Protocol) -> efi::Status {
    interrupts::disable_interrupts();

    efi::Status::SUCCESS
}

extern "efiapi" fn get_interrupt_state(this: *const Protocol, state: *mut bool) -> efi::Status {
    if state.is_null() {
        return efi::Status::INVALID_PARAMETER;
    }
    interrupts::get_interrupt_state()
        .map(|interrupt_state| {
            // SAFETY: caller must ensure that state is a valid pointer. It is null-checked above.
            unsafe {
                state.write_unaligned(interrupt_state);
            }
            efi::Status::SUCCESS
        })
        .unwrap_or_else(|err| err.into())
}

extern "efiapi" fn init(this: *const Protocol, init_type: CpuInitType) -> efi::Status {
    let Some(impl_ref) = get_impl_ref(this) else {
        return efi::Status::INVALID_PARAMETER;
    };

    let result = impl_ref.cpu.init(init_type);

    result.map(|_| efi::Status::SUCCESS).unwrap_or_else(|err| err.into())
}

extern "efiapi" fn register_interrupt_handler(
    this: *const Protocol,
    interrupt_type: isize,
    interrupt_handler: InterruptHandler,
) -> efi::Status {
    let Some(impl_ref) = get_impl_ref(this) else {
        return efi::Status::INVALID_PARAMETER;
    };
    let interrupt_manager = &impl_ref.interrupt_manager;

    let const_fn_ptr = interrupt_handler as *const ();
    let result = if const_fn_ptr.is_null() {
        interrupt_manager.unregister_exception_handler(interrupt_type as ExceptionType)
    } else {
        interrupt_manager
            .register_exception_handler(interrupt_type as ExceptionType, HandlerType::UefiRoutine(interrupt_handler))
    };

    match result {
        Ok(()) => efi::Status::SUCCESS,
        Err(err) => err.into(),
    }
}

extern "efiapi" fn get_timer_value(
    this: *const Protocol,
    timer_index: u32,
    timer_value: *mut u64,
    timer_period: *mut u64,
) -> efi::Status {
    if timer_value.is_null() || timer_period.is_null() {
        return efi::Status::INVALID_PARAMETER;
    }
    let Some(impl_ref) = get_impl_ref(this) else {
        return efi::Status::INVALID_PARAMETER;
    };

    let result = impl_ref.cpu.get_timer_value(timer_index);

    match result {
        Ok((value, period)) => {
            // SAFETY: caller must ensure that timer_value and timer_period are valid pointers. They are null-checked above.
            unsafe {
                timer_value.write_unaligned(value);
                timer_period.write_unaligned(period);
            }
            efi::Status::SUCCESS
        }
        Err(err) => err.into(),
    }
}

extern "efiapi" fn set_memory_attributes(
    _this: *const Protocol,
    base_address: efi::PhysicalAddress,
    length: u64,
    attributes: u64,
) -> efi::Status {
    match dxe_services::core_set_memory_space_attributes(base_address, length, attributes) {
        Ok(_) => efi::Status::SUCCESS,
        Err(status) => status.into(),
    }
}

impl EfiCpuArchProtocolImpl {
    fn new(cpu: Service<dyn Cpu>, interrupt_manager: Service<dyn InterruptManager>) -> Self {
        Self {
            protocol: Protocol {
                flush_data_cache,
                enable_interrupt,
                disable_interrupt,
                get_interrupt_state,
                init,
                register_interrupt_handler,
                get_timer_value,
                set_memory_attributes,
                number_of_timers: 0,
                dma_buffer_alignment: cpu.cache_writeback_granule(),
            },

            // private data
            cpu,
            interrupt_manager,
        }
    }
}

/// This component installs the cpu arch protocol
#[derive(Default)]
pub(crate) struct CpuArchProtocolInstaller;

#[component]
impl CpuArchProtocolInstaller {
    fn entry_point(
        self,
        cpu: Service<dyn Cpu>,
        interrupt_manager: Service<dyn InterruptManager>,
        bs: StandardBootServices,
    ) -> Result<()> {
        let protocol = EfiCpuArchProtocolImpl::new(cpu, interrupt_manager);

        // Convert the protocol to a raw pointer and store it in to protocol DB
        let interface = Box::leak(Box::new(protocol));

        bs.install_protocol_interface(None, interface)
            .inspect_err(|_| log::error!("Failed to install EFI_CPU_ARCH_PROTOCOL"))?;
        log::info!("installed EFI_CPU_ARCH_PROTOCOL_GUID");

        Ok(())
    }
}

#[cfg(test)]
#[coverage(off)]
mod tests {
    use crate::test_support;

    use super::*;

    use mockall::{mock, predicate::*};
    use patina::pi::protocols::cpu_arch::{EfiExceptionType, EfiSystemContext};

    mock! {
        EfiCpuInit {}
        impl Cpu for EfiCpuInit {
            fn flush_data_cache(
                &self,
                start: efi::PhysicalAddress,
                length: u64,
                flush_type: CpuFlushType,
            ) -> Result<()>;
            fn init(&self, init_type: CpuInitType) -> Result<()>;
            fn get_timer_value(&self, timer_index: u32) -> Result<(u64, u64)>;
            fn cache_writeback_granule(&self) -> u32;
        }
    }

    mock! {
        InterruptManager {}
        impl InterruptManager for InterruptManager {
            fn register_exception_handler(
                &self,
                interrupt_type: ExceptionType,
                handler: HandlerType,
            ) -> Result<()>;
            fn unregister_exception_handler(&self, interrupt_type: ExceptionType) -> Result<()>;
        }
    }

    fn with_locked_state<F: Fn() + std::panic::RefUnwindSafe>(f: F) {
        crate::test_support::with_global_lock(|| {
            test_support::init_test_logger();
            f();
        })
        .unwrap();
    }

    #[test]
    fn test_flush_data_cache() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            cpu_init.expect_flush_data_cache().with(eq(0), eq(0), always()).returning(|_, _, _| Ok(()));
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));

            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));

            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let status = flush_data_cache(&protocol.protocol, 0, 0, CpuFlushType::EfiCpuFlushTypeWriteBackInvalidate);
            assert_eq!(status, efi::Status::SUCCESS);

            // Verify the case when `this` is null.
            let status = flush_data_cache(core::ptr::null(), 0, 0, CpuFlushType::EfiCpuFlushTypeWriteBackInvalidate);
            assert_eq!(status, efi::Status::INVALID_PARAMETER);
        });
    }

    #[test]
    fn test_enable_interrupt() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));
            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let status = enable_interrupt(&protocol.protocol);
            assert_eq!(status, efi::Status::SUCCESS);
        });
    }

    #[test]
    fn test_disable_interrupt() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));
            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let status = disable_interrupt(&protocol.protocol);
            assert_eq!(status, efi::Status::SUCCESS);
        });
    }

    #[test]
    fn test_get_interrupt_state() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));
            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let mut state = false;
            let status = get_interrupt_state(&protocol.protocol, &mut state as *mut bool);
            assert_eq!(status, efi::Status::SUCCESS);
        });
    }

    #[test]
    fn test_init() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_init().with(always()).returning(|_| Ok(()));
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));

            let mut im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));

            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let status = init(&protocol.protocol, CpuInitType::EfiCpuInit);
            assert_eq!(status, efi::Status::SUCCESS);

            // Verify the case when `this` is null.
            let status = init(core::ptr::null(), CpuInitType::EfiCpuInit);
            assert_eq!(status, efi::Status::INVALID_PARAMETER);
        });
    }

    extern "efiapi" fn mock_interrupt_handler(_type: EfiExceptionType, _context: EfiSystemContext) {}

    #[test]
    fn test_register_interrupt_handler() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));

            let mut interrupt_manager = MockInterruptManager::new();
            interrupt_manager
                .expect_register_exception_handler()
                .with(eq(ExceptionType::from(0_usize)), always())
                .returning(|_, _| Ok(()));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(interrupt_manager));

            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let status = register_interrupt_handler(&protocol.protocol, 0, mock_interrupt_handler);
            assert_eq!(status, efi::Status::SUCCESS);

            // Verify the case when `this` is null.
            let status = register_interrupt_handler(core::ptr::null(), 0, mock_interrupt_handler);
            assert_eq!(status, efi::Status::INVALID_PARAMETER);
        });
    }

    #[test]
    fn test_get_timer_value() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            cpu_init.expect_get_timer_value().with(eq(0)).returning(|_| Ok((0, 0)));
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));

            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));

            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let mut timer_value: u64 = 0;
            let mut timer_period: u64 = 0;
            let status =
                get_timer_value(&protocol.protocol, 0, &mut timer_value as *mut _, &mut timer_period as *mut _);
            assert_eq!(status, efi::Status::SUCCESS);

            // Verify the case when `this` is null.
            let status = get_timer_value(core::ptr::null(), 0, &mut timer_value as *mut _, &mut timer_period as *mut _);
            assert_eq!(status, efi::Status::INVALID_PARAMETER);

            // Null out-parameters should also be rejected.
            let status = get_timer_value(&protocol.protocol, 0, core::ptr::null_mut(), &mut timer_period as *mut _);
            assert_eq!(status, efi::Status::INVALID_PARAMETER);
            let status = get_timer_value(&protocol.protocol, 0, &mut timer_value as *mut _, core::ptr::null_mut());
            assert_eq!(status, efi::Status::INVALID_PARAMETER);
        });
    }

    // Tests for DxeCpu delegation
    #[test]
    fn test_dxe_cpu_flush_data_cache_delegates() {
        with_locked_state(|| {
            let dxe_cpu = DxeCpu(EfiCpu::default());
            let result = dxe_cpu.flush_data_cache(0x1000, 0x100, CpuFlushType::EfiCpuFlushTypeWriteBackInvalidate);
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_dxe_cpu_init_delegates() {
        with_locked_state(|| {
            let dxe_cpu = DxeCpu(EfiCpu::default());
            let result = dxe_cpu.init(CpuInitType::EfiCpuInit);
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_dxe_cpu_get_timer_value_delegates() {
        with_locked_state(|| {
            let dxe_cpu = DxeCpu(EfiCpu::default());
            let result = dxe_cpu.get_timer_value(0);
            assert_eq!(result.unwrap(), (0, 0));
        });
    }

    #[test]
    fn test_dxe_cpu_cache_writeback_granule_delegates() {
        with_locked_state(|| {
            let dxe_cpu = DxeCpu(EfiCpu::default());
            let granule = dxe_cpu.cache_writeback_granule();
            assert!(granule == 64u32);
        });
    }

    // Tests for DxeInterruptManager delegation
    #[test]
    fn test_dxe_interrupt_manager_register_then_unregister_delegates() {
        with_locked_state(|| {
            let dxe_interrupt_manager = DxeInterruptManager(Interrupts::default());

            // Register first
            let result = dxe_interrupt_manager.register_exception_handler(
                ExceptionType::from(0_usize),
                HandlerType::UefiRoutine(mock_interrupt_handler),
            );
            assert!(result.is_ok());

            // Then unregister
            let result = dxe_interrupt_manager.unregister_exception_handler(ExceptionType::from(0_usize));
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_dxe_interrupt_manager_unregister_then_register_delegates() {
        with_locked_state(|| {
            let dxe_interrupt_manager = DxeInterruptManager(Interrupts::default());
            let result = dxe_interrupt_manager.unregister_exception_handler(ExceptionType::from(0_usize));
            // Expecting an error because there is no handler registered yet, but the method should still be callable.
            assert!(result.is_err());

            let result = dxe_interrupt_manager.register_exception_handler(
                ExceptionType::from(0_usize),
                HandlerType::UefiRoutine(mock_interrupt_handler),
            );
            assert!(result.is_ok());

            // Now the unregister should succeed
            let result = dxe_interrupt_manager.unregister_exception_handler(ExceptionType::from(0_usize));
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_get_impl_ref_null_returns_none() {
        assert!(get_impl_ref(core::ptr::null()).is_none());
    }

    #[test]
    fn test_get_impl_ref_returns_some_for_valid_pointer() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));
            let protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let this = &raw const protocol.protocol;
            let impl_ref = get_impl_ref(this).expect("non-null pointer should yield Some");
            assert_eq!(&raw const impl_ref.protocol, this);
        });
    }

    #[test]
    fn test_get_impl_ref_mut_null_returns_none() {
        assert!(get_impl_ref_mut(core::ptr::null_mut()).is_none());
    }

    #[test]
    fn test_get_impl_ref_mut_returns_some_for_valid_pointer() {
        with_locked_state(|| {
            let mut cpu_init = MockEfiCpuInit::new();
            cpu_init.expect_cache_writeback_granule().return_const(64_u32);
            let cpu: Service<dyn Cpu> = Service::mock(Box::new(cpu_init));
            let im: Service<dyn InterruptManager> = Service::mock(Box::new(MockInterruptManager::new()));
            let mut protocol = EfiCpuArchProtocolImpl::new(cpu, im);

            let this = &raw mut protocol.protocol;
            let impl_ref = get_impl_ref_mut(this).expect("non-null pointer should yield Some");
            assert_eq!(&raw const impl_ref.protocol, this.cast_const());
        });
    }
}
