cpu-temp 0.1.0

use std::num::NonZeroU16;

use raw_cpuid::{CpuId as RawCpuId, TopologyType};

#[cfg(target_os = "windows")]
use windows::Win32::{Foundation::*, System::SystemInformation::*};

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Vendor {
    Intel,
    Amd,
    Unknown,
}

pub struct CpuInfo {
    vendor: Vendor,
    family: u32,
    model: u32,
    stepping: u32,
    core_count: Option<NonZeroU16>,
}

impl CpuInfo {
    pub fn new() -> Option<Self> {
        let cpuid = RawCpuId::new();

        // 获取厂商信息
        let vendor = match cpuid.get_vendor_info() {
            Some(info) => match info.as_str() {
                "GenuineIntel" => Vendor::Intel,
                "AuthenticAMD" => Vendor::Amd,
                _ => Vendor::Unknown,
            },
            None => Vendor::Unknown,
        };

        // 获取基础CPU信息
        let (family, model, stepping) = match cpuid.get_feature_info() {
            Some(info) => (
                info.family_id() as u32,
                info.model_id() as u32,
                info.stepping_id() as u32,
            ),
            None => (0, 0, 0),
        };

        let core_count = cpuid
            .get_extended_topology_info()
            .and_then(|mut info| info.find(|level| level.level_type() == TopologyType::Core))
            .and_then(|level| NonZeroU16::new(level.processors()));

        Some(CpuInfo {
            vendor,
            family,
            model,
            stepping,
            core_count,
        })
    }

    pub fn get_vendor(&self) -> Vendor {
        self.vendor
    }

    pub fn get_family(&self) -> u32 {
        self.family
    }

    pub fn get_model(&self) -> u32 {
        self.model
    }

    pub fn get_stepping(&self) -> u32 {
        self.stepping
    }

    pub fn get_core_count(&self) -> Option<NonZeroU16> {
        self.core_count
    }

    /// 检查CPU是否支持温度传感器
    pub fn has_temperature_sensor(&self) -> bool {
        let cpuid = RawCpuId::new();
        if cpuid.get_feature_info().is_some() {
            // 检查是否支持数字温度传感器
            // 通过CPUID功能06H的EAX返回值的第0位检查
            if let Some(thermal_info) = cpuid.get_thermal_power_info() {
                return thermal_info.has_dts(); // 检查是否支持数字温度传感器
            }
        }
        false
    }

    /// 获取CPU名称
    pub fn get_name(&self) -> String {
        let cpuid = RawCpuId::new();
        if let Some(info) = cpuid.get_processor_brand_string() {
            info.as_str().to_string()
        } else {
            "Unknown CPU".to_string()
        }
    }

    /// 获取物理核心对应的逻辑CPU核心映射
    /// 返回一个向量，每个元素代表一个物理核心，包含该核心对应的逻辑CPU核心列表
    pub fn get_core_cpu_mapping() -> anyhow::Result<Vec<Vec<usize>>> {
        #[cfg(target_os = "windows")]
        {
            Self::get_core_cpu_mapping_windows()
        }
        #[cfg(target_os = "linux")]
        {
            let cpu_info = CpuInfo::new().ok_or_else(|| anyhow::anyhow!("Failed to get CPU info"))?;
            cpu_info.get_core_cpu_mapping_linux()
        }
        #[cfg(not(any(target_os = "windows", target_os = "linux")))]
        {
            // 默认实现，返回空向量
            anyhow::bail!("not implementated");
        }
    }

    /// Windows实现：获取物理核心对应的逻辑CPU核心映射
    #[cfg(target_os = "windows")]
    fn get_core_cpu_mapping_windows() -> anyhow::Result<Vec<Vec<usize>>> {
        use std::mem;

        let mut mappings = Vec::new();

        // 获取逻辑处理器信息所需的缓冲区大小
        let mut buffer_length: u32 = 0;
        let result =
            unsafe { GetLogicalProcessorInformationEx(RelationAll, None, &mut buffer_length as _) };

        match result {
            Ok(()) => return Ok(mappings),
            Err(e) => {
                if unsafe { GetLastError() } != ERROR_INSUFFICIENT_BUFFER {
                    anyhow::bail!(
                        "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
                        e.code().0,
                        e.message()
                    )
                }
            }
        }

        // 分配缓冲区
        let mut buffer: Vec<u8> = vec![0; buffer_length as usize];

        // 获取逻辑处理器信息
        let result = unsafe {
            GetLogicalProcessorInformationEx(
                RelationAll,
                Some(buffer.as_mut_ptr() as _),
                &mut buffer_length as _,
            )
        };

        if let Err(e) = result {
            anyhow::bail!(
                "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
                e.code().0,
                e.message()
            )
        }

        // 解析处理器关系信息
        let mut offset: usize = 0;
        while offset < buffer.len() {
            use anyhow::Context;
            let size = mem::size_of::<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>();
            let cur_buffer = buffer
                .get(offset..offset + size)
                .context("slice info error")?;
            let info: SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX =
                unsafe { std::ptr::read_unaligned(cur_buffer.as_ptr() as _) };
            if info.Relationship == RelationProcessorCore {
                let mut logical_processors = Vec::new();

                // 获取处理器关系信息
                let processor_info = unsafe { &info.Anonymous.Processor };

                let group_mask_offset = offset
                    + mem::offset_of!(
                        SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX,
                        Anonymous.Processor.GroupMask
                    );
                let group_mask_size = processor_info.GroupCount as usize
                    * mem::size_of_val(&processor_info.GroupMask[0]);
                let mask_buffer = buffer
                    .get(group_mask_offset..group_mask_offset + group_mask_size)
                    .context("slice group mask error")?;

                // 遍历组掩码
                for i in 0..processor_info.GroupCount as usize {
                    let group_mask = unsafe {
                        mask_buffer
                            .as_ptr()
                            .cast::<GROUP_AFFINITY>()
                            .add(i)
                            .read_unaligned()
                    };
                    let mut mask = group_mask.Mask;

                    // 解析掩码中的逻辑处理器
                    while mask != 0 {
                        // 获取最低位 1 的索引
                        let bit = mask.trailing_zeros();
                        logical_processors.push(i * 64 + bit as usize);

                        // 清除最低位的 1
                        mask &= mask - 1;
                    }
                }

                mappings.push(logical_processors);
            }

            // 移动到下一个条目
            offset += info.Size as usize;
        }

        Ok(mappings)
    }

    /// Linux实现：获取物理核心对应的逻辑CPU核心映射
    #[cfg(target_os = "linux")]
    fn get_core_cpu_mapping_linux(&self) -> anyhow::Result<Vec<Vec<usize>>> {
        use std::collections::BTreeMap;
        use std::fs;

        use anyhow::Context as _;

        const CORE_ID_PATTERN: &str = "/sys/devices/system/cpu/cpu[0-9]*/topology/core_id";

        let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
        for entry in glob::glob(CORE_ID_PATTERN).context("Failed to read glob pattern")? {
            if let Ok(path) = entry {
                // 获取目录名中的 cpu 号，例如从 ".../cpu3/topology/core_id" 提取 3
                let cpu_id = path
                    .ancestors()
                    .nth(2)
                    .and_then(|p| p.file_name())
                    .and_then(|s| s.to_str())
                    .and_then(|s| s.trim_start_matches("cpu").parse::<usize>().ok());

                // 读取 core_id 文件内容
                let core_id = fs::read_to_string(&path)
                    .ok()
                    .and_then(|s| s.trim().parse::<usize>().ok());

                if let (Some(cpu), Some(core)) = (cpu_id, core_id) {
                    groups.entry(core).or_default().push(cpu);
                }
            }
        }

        // 将分组后的结果转为 Vec<Vec<usize>>
        Ok(groups.into_values().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_core_cpu_mapping() {
        // use windows::Win32::System::Threading::*;
        // let mut groupcount = 0u16;
        // let handle = unsafe { GetCurrentProcess() };
        // let _ =
        //     unsafe { GetProcessGroupAffinity(handle, &mut groupcount as _, Default::default()) };
        // let grouparray: Vec<u16> = vec![0; groupcount as usize];
        // let res = unsafe {
        //     GetProcessGroupAffinity(handle, &mut groupcount as _, grouparray.as_ptr() as _)
        // };
        // if !res.as_bool() {
        //     panic!("failed: {}", unsafe {
        //         &GetLastError().to_hresult().message()
        //     });
        // }
        let res = CpuInfo::get_core_cpu_mapping().unwrap();
        println!("{:?}", res);
    }
}