use nvml_wrapper::Nvml;

use crate::dtos::gpu::GpuInfo;

pub fn get_gpu_info() -> Vec<GpuInfo> {
	// Initialize NVML
	let nvml = Nvml::init().expect("Cannot init NVML");

	// Get the count of available GPUs
	let device_count = nvml.device_count().expect("Cannot get GPU's count");

	// Create a vector to store the GPU information
	let mut gpu_info_list = vec![];

	// Iterate through each GPU
	for i in 0..device_count {
		 // Get the device by index
		 let device = nvml.device_by_index(i).unwrap();

		 // Collect GPU information and push it to the vector
		 gpu_info_list.push(GpuInfo {
			  name: device.name().expect("Cannot get GPU name"),
			  virtual_ram: device
					.memory_info()
					.expect("Cannot load GPU memmory info")
					.total,
		 })
	}

	// Return the vector of GPU information
	gpu_info_list
}
