use amd_smi_wrapper::{
    AmdInterface,
    handles::{ProcessorHandle, SocketHandle},
};
use std::collections::HashMap;

use super::features::OptionalFeatures;

/// SAFETY: The amd libary is thread-safe and returns pointers to a safe global state, which we can pass to other threads.
unsafe impl<H: ProcessorHandle> Send for ManagedDevice<H> {}

/// Detected AMD GPU devices via AMDSMI.
pub struct AmdGpuDevices<H: ProcessorHandle> {
    /// Counter of detection errors on AMD GPU device.
    pub failure_count: usize,
    /// Set of parameters that defines an AMD GPU device.
    pub devices: Vec<ManagedDevice<H>>,
}

/// An AMD GPU device that has been probed for available features.
pub struct ManagedDevice<H: ProcessorHandle> {
    /// A pointer to the device.
    pub handle: H,
    /// Status of the various features available or not on a device.
    pub features: OptionalFeatures,
    /// PCI bus ID of the device.
    pub bus_id: String,
}

/// Statistics about the device detection.
pub struct DetectionStats {
    /// Detected AMD GPU device.
    pub found_devices: usize,
    /// Detected working features on the AMD GPU device.
    pub working_devices: usize,
}

impl<H: ProcessorHandle> AmdGpuDevices<H> {
    /// Detects all AMD GPUs and returns an AmdGpuDevices object.
    pub fn detect<A>(amdsmi: &A, skip_failed_devices: bool) -> anyhow::Result<Self>
    where
        A: AmdInterface<SocketHandle: SocketHandle<ProcessorHandle = H>>,
    {
        let mut devices = HashMap::new();
        let mut failure_count = 0;

        // Iterate over sockets
        for socket in amdsmi.socket_handles()? {
            // Iterate over processor handles
            for processor in socket.processor_handles()? {
                // Detect available features
                match OptionalFeatures::with_detected_features(&processor) {
                    Ok((_, features)) => {
                        let bus_id = processor.device_uuid()?.to_string();

                        if features.has_any() {
                            devices.insert(
                                bus_id.clone(),
                                ManagedDevice {
                                    handle: processor,
                                    features,
                                    bus_id,
                                },
                            );
                        } else {
                            log::warn!("Skipping GPU device because it supports no useful feature.");
                            failure_count += 1;
                        }
                    }
                    Err(e) => {
                        if skip_failed_devices {
                            failure_count += 1;
                            log::warn!("Skipping GPU device because of error: {e:?}");
                        } else {
                            return Err(e.into());
                        }
                    }
                }
            }
        }

        let mut devices: Vec<_> = devices.into_values().collect();
        devices.sort_by_key(|d| d.bus_id.clone());

        Ok(AmdGpuDevices { devices, failure_count })
    }

    /// Gets statistics about device detection.
    pub fn detection_stats(&self) -> DetectionStats {
        let n_failed = self.failure_count;
        let n_ok = self.devices.len();
        let n_found = n_ok + n_failed;

        DetectionStats {
            found_devices: n_found,
            working_devices: n_ok,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::tests::mocks::{
        MOCK_ACTIVITY, MOCK_ENERGY, MOCK_MEMORY, MOCK_POWER, MOCK_TEMPERATURE, MOCK_UUID, MOCK_VOLTAGE, mock_asic_info,
        mock_process,
    };
    use amd_smi_wrapper::{
        MockAmdInterface,
        error::AmdError,
        handles::{MockProcessorHandle, MockSocketHandle},
        metrics::AmdTemperatureMetric,
    };
    use amd_smi_wrapper_sys::versions::stable::{AMDSMI_STATUS_ADDRESS_FAULT, AMDSMI_STATUS_INVAL};
    use log::LevelFilter::Warn;
    use std::{
        io::Write,
        sync::{Arc, Mutex},
    };

    // Test `detection_stats` function with no GPUs detected
    #[test]
    fn test_detection_stats_no_devices() {
        let devices: AmdGpuDevices<MockProcessorHandle> = AmdGpuDevices {
            devices: vec![],
            failure_count: 0,
        };

        let stats = devices.detection_stats();
        assert_eq!(stats.found_devices, 0);
        assert_eq!(stats.working_devices, 0);
    }

    // Test `detection_stats` function with GPUs detected but no working device
    #[test]
    fn test_detection_stats_failed() {
        let devices: AmdGpuDevices<MockProcessorHandle> = AmdGpuDevices {
            devices: vec![],
            failure_count: 5,
        };

        let stats = devices.detection_stats();
        assert_eq!(stats.found_devices, 5);
        assert_eq!(stats.working_devices, 0);
    }

    // Test `detect` function in success case with valid GPUs and metrics
    #[test]
    fn test_detect_success() {
        let mut mock_init = MockAmdInterface::new();
        let mut mock_socket = MockSocketHandle::new();
        let mut mock_processor = MockProcessorHandle::new();

        mock_processor
            .expect_device_uuid()
            .returning(|| Ok(MOCK_UUID.to_owned()));

        mock_processor.expect_device_activity().returning(|| Ok(MOCK_ACTIVITY));

        mock_processor
            .expect_device_asic_info()
            .returning(|| Ok(mock_asic_info()));

        mock_processor
            .expect_device_energy_consumption()
            .returning(|| Ok(MOCK_ENERGY[0]));

        mock_processor.expect_device_power_info().returning(|| Ok(MOCK_POWER));

        mock_processor
            .expect_is_power_managment_enabled()
            .returning(|| Ok(true));
        mock_processor
            .expect_device_process_list()
            .returning(|| Ok(vec![mock_process()]));
        mock_processor
            .expect_device_voltage()
            .returning(|_, _| Ok(MOCK_VOLTAGE));

        mock_processor.expect_device_memory_usage().returning(|mem_type| {
            MOCK_MEMORY
                .iter()
                .find(|(t, _)| *t == mem_type)
                .map(|(_, v)| Ok(*v))
                .unwrap_or(Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)))
        });

        mock_processor.expect_device_temperature().returning(|sensor, metric| {
            if metric != AmdTemperatureMetric::AMDSMI_TEMP_CURRENT {
                return Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL));
            }
            MOCK_TEMPERATURE
                .iter()
                .find(|(s, _)| *s == sensor)
                .map(|(_, v)| Ok(*v))
                .unwrap_or(Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)))
        });

        mock_socket
            .expect_processor_handles()
            .return_once(move || Ok(vec![mock_processor]));

        mock_init
            .expect_socket_handles()
            .return_once(move || Ok(vec![mock_socket]));

        let res = AmdGpuDevices::detect(&mock_init, false).expect("should work");
        assert_eq!(res.failure_count, 0);
    }

    // Test `detect` function for a GPU with no features available
    #[test]
    fn test_detect_error_skipped() {
        let mut mock_init = MockAmdInterface::new();
        let mut mock_socket = MockSocketHandle::new();
        let mut mock_processor = MockProcessorHandle::new();

        mock_processor
            .expect_device_uuid()
            .returning(|| Ok(MOCK_UUID.to_owned()));

        mock_processor
            .expect_device_activity()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_asic_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_energy_consumption()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_power_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_is_power_managment_enabled()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_process_list()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_voltage()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_memory_usage()
            .returning(|_| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));
        mock_processor
            .expect_device_temperature()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_INVAL)));

        mock_socket
            .expect_processor_handles()
            .return_once(move || Ok(vec![mock_processor]));

        mock_init
            .expect_socket_handles()
            .return_once(move || Ok(vec![mock_socket]));

        let res = AmdGpuDevices::detect(&mock_init, true).expect("should skip errors");

        assert_eq!(res.failure_count, 1);
        assert!(res.devices.is_empty());
    }

    // Test `detect` function with not successfully skipped features of a GPU
    #[test]
    fn test_detect_error_not_skipped() {
        let mut mock_init = MockAmdInterface::new();
        let mut mock_socket = MockSocketHandle::new();
        let mut mock_processor = MockProcessorHandle::new();

        mock_processor
            .expect_device_uuid()
            .returning(|| Ok(MOCK_UUID.to_owned()));

        mock_processor
            .expect_device_activity()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_asic_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_energy_consumption()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_power_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_is_power_managment_enabled()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_process_list()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_voltage()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_memory_usage()
            .returning(|_| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_temperature()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));

        mock_socket
            .expect_processor_handles()
            .return_once(move || Ok(vec![mock_processor]));

        mock_init
            .expect_socket_handles()
            .return_once(move || Ok(vec![mock_socket]));

        let res = AmdGpuDevices::detect(&mock_init, false);
        assert!(res.is_err());
    }

    // Test `detect` function with error skipped and logged
    #[test]
    fn test_detect_error_skipped_logged() {
        let log = Arc::new(Mutex::new(Vec::new()));
        let value = Arc::clone(&log);

        let _ = env_logger::builder()
            .is_test(true)
            .format(move |buf, record| {
                value.lock().unwrap().push(format!("{}", record.args()));
                writeln!(buf, "{}", record.args())
            })
            .filter_level(Warn)
            .try_init();

        let mut mock_init = MockAmdInterface::new();
        let mut mock_socket = MockSocketHandle::new();
        let mut mock_processor = MockProcessorHandle::new();

        mock_processor
            .expect_device_uuid()
            .returning(|| Ok(MOCK_UUID.to_owned()));

        mock_processor
            .expect_device_activity()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_asic_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_energy_consumption()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_power_info()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_is_power_managment_enabled()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_process_list()
            .returning(|| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_voltage()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_memory_usage()
            .returning(|_| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));
        mock_processor
            .expect_device_temperature()
            .returning(|_, _| Err(AmdError::from_status_code(AMDSMI_STATUS_ADDRESS_FAULT)));

        mock_socket
            .expect_processor_handles()
            .return_once(move || Ok(vec![mock_processor]));

        mock_init
            .expect_socket_handles()
            .return_once(move || Ok(vec![mock_socket]));

        let res = AmdGpuDevices::detect(&mock_init, true).expect("should skip the failed GPU");

        assert_eq!(res.failure_count, 1);
        assert!(res.devices.is_empty());
        assert!(
            log.lock()
                .unwrap()
                .iter()
                .any(|msg| msg.contains("Skipping GPU device because of error"))
        );
    }
}
