//! Handles used to manipulate the devices.

use std::{
    ffi::c_char,
    mem::MaybeUninit,
    ptr::{self},
};

use crate::{
    AmdSmi,
    error::AmdError,
    metrics::{
        asic_info::AmdAsicInfo, clock_info::AmdClockInfo, power_info::AmdPowerInfo,
        process_info::AmdProcessInfo, *,
    },
    utils::c_buffer_to_string,
};

use amd_smi_wrapper_sys::versions::stable::{self, amdsmi_processor_handle, amdsmi_socket_handle};

#[cfg(feature = "mock")]
use mockall::automock;

pub struct AmdSocketHandle {
    pub(crate) amdsmi: AmdSmi,
    pub(crate) inner: amdsmi_socket_handle,
}

pub struct AmdProcessorHandle {
    pub(crate) amdsmi: AmdSmi,
    pub(crate) inner: amdsmi_processor_handle,
}

/// Handle to a socket in the system.
#[cfg_attr(feature = "mock", automock(type ProcessorHandle=MockProcessorHandle;))]
pub trait SocketHandle {
    /// The type of processor handles returned by this socket.
    type ProcessorHandle: ProcessorHandle;

    /// Lists the processors associated to this socket.
    fn processor_handles(&self) -> Result<Vec<Self::ProcessorHandle>, AmdError>;
}

impl SocketHandle for AmdSocketHandle {
    type ProcessorHandle = AmdProcessorHandle;

    fn processor_handles(&self) -> Result<Vec<Self::ProcessorHandle>, AmdError> {
        let mut processor_count = 0;

        // Query the number of processor handles for the given socket.
        // SAFETY: According the AMD-SMI library documentation, passing `null_mut()` is safe which sets `processor_count` to the number of processors available for this socket.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_processor_handles(self.inner, &mut processor_count, ptr::null_mut())
        };
        self.amdsmi.check_status(result)?;

        // Allocate a vector of nulls.
        let mut processor_handles = vec![ptr::null_mut(); processor_count as usize];

        // Fill the buffer with processor handles.
        // SAFETY: `processor_handles.as_mut_ptr()` points to a memory block of sufficient size.
        //  According the AMD-SMI library documentation, the function writes at most `processor_count` handles ensuring no out-of-bounds access occurs.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_processor_handles(
                    self.inner,
                    &mut processor_count,
                    processor_handles.as_mut_ptr(),
                )
        };

        self.amdsmi.check_status(result)?;

        processor_handles.truncate(processor_count as usize);
        Ok(processor_handles
            .into_iter()
            .map(|s| AmdProcessorHandle {
                amdsmi: self.amdsmi.clone(),
                inner: s,
            })
            .collect())
    }
}

/// Handle to a processor in a [socket](SocketHandle).
#[cfg_attr(feature = "mock", automock)]
pub trait ProcessorHandle {
    /// Retrieves a [`AmdEngineUsage`] structure containing all data about GPU device activities.
    fn device_activity(&self) -> Result<AmdEngineUsage, AmdError>;

    /// Retrieves globals [`AmdAsicInfo`] information about a GPU device.
    fn device_asic_info(&self) -> Result<AmdAsicInfo, AmdError>;

    /// Retrieves metrics about the GPU clock.
    ///
    /// # Arguments
    ///
    /// - `clk_type`: Clock devices existing among [`AmdClkType`] on hardware.
    fn device_clock_info(&self, clk_type: AmdClkType) -> Result<AmdClockInfo, AmdError>;

    /// Retrieves a [`AmdEnergyConsumption`] structure containing data about energy consumption of the GPU device.
    fn device_energy_consumption(&self) -> Result<AmdEnergyConsumption, AmdError>;

    /// Retrieves the fan speed ratio.
    fn device_fan_speed(&self, sensor_index: u32) -> Result<f32, AmdError>;

    /// Retrieves the memory consumption of the GPU device.
    ///
    /// # Arguments
    ///
    /// - `mem_type`: Memory devices existing among [`AmdMemoryType`] on hardware.
    fn device_memory_usage(&self, mem_type: AmdMemoryType) -> Result<u64, AmdError>;

    /// Retrieves the PCI bus traffic used by the GPU device.
    fn device_pci_usage(&self) -> Result<AmdPciTraffic, AmdError>;

    /// Retrieves a [`AmdPowerInfo`] structure containing all data about GPU device power consumption.
    fn device_power_info(&self) -> Result<AmdPowerInfo, AmdError>;
    /// Retrieves the power management status accessability of the GPU device.
    fn is_power_managment_enabled(&self) -> Result<bool, AmdError>;

    /// Retrieves data about the processes that are currently running on the GPU device.
    fn device_process_list(&self) -> Result<Vec<AmdProcessInfo>, AmdError>;

    /// Retrieves the temperature of a given area of the GPU device.
    ///
    /// # Arguments
    ///
    /// - `sensor_type`: Thermal sensor [`AmdTemperatureType`] on AMD GPU hardware.
    /// - `metric`: Temperature type [`AmdTemperatureMetric`] analyzed (current, average...).
    fn device_temperature(
        &self,
        sensor_type: AmdTemperatureType,
        metric: AmdTemperatureMetric,
    ) -> Result<i64, AmdError>;

    /// Retrieves the UUID of the GPU device.
    fn device_uuid(&self) -> Result<String, AmdError>;

    /// Retrieves the voltage of a given area of the GPU device, in milliVolt.
    ///
    /// # Arguments
    ///
    /// - `sensor_type`: Voltage retrieved by a [`AmdVoltageType`] sensor on AMD GPU hardware.
    /// - `metric`: Voltage type [`AmdVoltageMetric`] analyzed (current, average...).
    fn device_voltage(
        &self,
        sensor_type: AmdVoltageType,
        metric: AmdVoltageMetric,
    ) -> Result<i64, AmdError>;
}

impl ProcessorHandle for AmdProcessorHandle {
    fn device_activity(&self) -> Result<AmdEngineUsage, AmdError> {
        // Allocate uninitialized memory for the structure and avoid reading uninitialized memory before the FFI call.
        let mut info = MaybeUninit::<stable::amdsmi_engine_usage_t>::uninit();

        // SAFETY: Pass a raw pointer to uninitialized memory to the FFI function.
        // According to AMD-SMI documentation, the function fully initializes the structure on success.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_activity(self.inner, info.as_mut_ptr())
        };

        self.amdsmi.check_status(result)?;

        // SAFETY: `assume_init()` is safe because the FFI call succeeded and fully initialized `info`.
        let info = unsafe { info.assume_init() };
        Ok(info.into())
    }

    fn device_asic_info(&self) -> Result<AmdAsicInfo, AmdError> {
        AmdAsicInfo::get(self)
    }

    fn device_clock_info(&self, clk_type: AmdClkType) -> Result<AmdClockInfo, AmdError> {
        AmdClockInfo::get(self, clk_type)
    }

    fn device_energy_consumption(&self) -> Result<AmdEnergyConsumption, AmdError> {
        let mut consumption = AmdEnergyConsumption::default();

        // SAFETY: Pass mutable pointers to the fields of `consumption` to the FFI function.
        // According to AMD-SMI documentation, the function writes all values on success and will not write beyond the memory locations provided.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi.shared.inner.lib_stable.amdsmi_get_energy_count(
                self.inner,
                &mut consumption.energy,
                &mut consumption.resolution,
                &mut consumption.timestamp,
            )
        };

        self.amdsmi.check_status(result)?;
        Ok(consumption)
    }

    fn device_fan_speed(&self, sensor_index: u32) -> Result<f32, AmdError> {
        let mut speed = 0;

        // SAFETY: Pass a mutable pointer to `speed` for the FFI function to write the current fan speed.
        // According to AMD-SMI documentation, the function writes a value between 0 and `AMDSMI_MAX_FAN_SPEED` to this pointer.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_fan_speed(self.inner, sensor_index, &mut speed)
        };

        self.amdsmi.check_status(result)?;
        Ok((speed as f32 * 100.0) / stable::AMDSMI_MAX_FAN_SPEED as f32)
    }

    fn device_memory_usage(&self, mem_type: AmdMemoryType) -> Result<u64, AmdError> {
        let mut used = 0;

        // SAFETY: Pass a mutable pointer to `used` for the FFI function to write the memory usage.
        // According to AMD-SMI documentation, the function will write a valid value on success and will not write outside the provided memory location.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_memory_usage(self.inner, mem_type, &mut used)
        };

        self.amdsmi.check_status(result)?;
        Ok(used)
    }

    fn device_pci_usage(&self) -> Result<AmdPciTraffic, AmdError> {
        let mut usage = AmdPciTraffic {
            sent: 0,
            received: 0,
            max_pkt_sz: 0,
        };

        // SAFETY: Pass mutable pointers to the fields of `usage` to the FFI function.
        // According to AMD-SMI documentation, the function writes all values on success or ignored them.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_pci_throughput(
                    self.inner,
                    &mut usage.sent,
                    &mut usage.received,
                    &mut usage.max_pkt_sz,
                )
        };

        self.amdsmi.check_status(result)?;
        Ok(usage)
    }

    fn device_power_info(&self) -> Result<AmdPowerInfo, AmdError> {
        AmdPowerInfo::get(self)
    }

    fn is_power_managment_enabled(&self) -> Result<bool, AmdError> {
        let mut enabled = false;

        // SAFETY: Pass a mutable pointer to `enabled` for the FFI function to write the power management status.
        // According to AMD-SMI documentation, the function will write a valid boolean value on success.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_is_gpu_power_management_enabled(self.inner, &mut enabled)
        };

        self.amdsmi.check_status(result)?;
        Ok(enabled)
    }

    fn device_temperature(
        &self,
        sensor_type: AmdTemperatureType,
        metric: AmdTemperatureMetric,
    ) -> Result<i64, AmdError> {
        let mut temperature = 0;

        // SAFETY: Pass a mutable pointer to `temperature` for the FFI function to write the temperature value.
        // According to AMD-SMI documentation, the function writes the value to this pointer.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi.shared.inner.lib_stable.amdsmi_get_temp_metric(
                self.inner,
                sensor_type,
                metric,
                &mut temperature,
            )
        };

        self.amdsmi.check_status(result)?;
        Ok(temperature)
    }

    fn device_voltage(
        &self,
        sensor_type: AmdVoltageType,
        metric: AmdVoltageMetric,
    ) -> Result<i64, AmdError> {
        let mut voltage = 0;

        // SAFETY: Pass a non-null mutable pointer to `voltage` for the FFI function to write the voltage value.
        // According to AMD-SMI documentation, the function writes the value to this pointer.
        // The value is only read after confirming that the return status is SUCCESS.
        // The `SUCCESS` return code `amdsmi_status_t` is checked before using the data.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_volt_metric(self.inner, sensor_type, metric, &mut voltage)
        };

        self.amdsmi.check_status(result)?;
        Ok(voltage)
    }

    fn device_process_list(&self) -> Result<Vec<AmdProcessInfo>, AmdError> {
        AmdProcessInfo::list(self)
    }

    fn device_uuid(&self) -> Result<String, AmdError> {
        let mut uuid_buffer = [0 as c_char; stable::AMDSMI_GPU_UUID_SIZE as usize];
        let mut uuid_length = stable::AMDSMI_GPU_UUID_SIZE;

        // SAFETY: According to AMD-SMI documentation, the function will not write beyond `uuid_length`.
        // `uuid_length` must be initialized to the buffer size, and the function will update it with the actual length.
        let result = unsafe {
            self.amdsmi
                .shared
                .inner
                .lib_stable
                .amdsmi_get_gpu_device_uuid(self.inner, &mut uuid_length, uuid_buffer.as_mut_ptr())
        };

        self.amdsmi.check_status(result)?;

        let uuid = c_buffer_to_string(uuid_buffer.as_slice());
        Ok(uuid)
    }
}
