zenyx-engine-telemetry/engine/src/metadata.rs

use std::fmt;
use std::str::FromStr;
use sysinfo::{CpuRefreshKind, RefreshKind, System};
use raw_cpuid::CpuId;
use wgpu::DeviceType;
use std::collections::HashSet;

mod build_info {
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Memory {
    bytes: u64,
}

impl Memory {
    pub const fn from_bytes(bytes: u64) -> Self {
        Self { bytes }
    }

    pub const fn from_kb(kb: u64) -> Self {
        Self { bytes: kb * 1024 }
    }

    pub const fn from_mb(mb: u64) -> Self {
        Self { bytes: mb * 1024 * 1024 }
    }

    pub const fn from_gb(gb: u64) -> Self {
        Self { bytes: gb * 1024 * 1024 * 1024 }
    }

    pub const fn as_bytes(&self) -> u64 {
        self.bytes
    }

    pub const fn as_kb(&self) -> u64 {
        self.bytes / 1024
    }

    pub const fn as_mb(&self) -> u64 {
        self.bytes / (1024 * 1024)
    }

    pub const fn as_gb(&self) -> u64 {
        self.bytes / (1024 * 1024 * 1024)
    }

    pub fn format_human(&self) -> String {
        const UNITS: [&str; 4] = ["B", "KB", "MB", "GB"];
        let mut size = self.bytes as f64;
        let mut unit_index = 0;

        while size >= 1024.0 && unit_index < UNITS.len() - 1 {
            size /= 1024.0;
            unit_index += 1;
        }

        format!("{:.2} {}", size, UNITS[unit_index])
    }

    pub fn verbose_info(&self) -> String {
        format!(
            "{} ({} bytes, {} KB, {} MB, {} GB)",
            self.format_human(),
            self.as_bytes(),
            self.as_kb(),
            self.as_mb(),
            self.as_gb()
        )
    }
}

impl fmt::Display for Memory {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.format_human())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CPUBrand {
    Intel,
    AMD,
    SnapDragon,
    Apple,
    Other(String),
}

impl From<&str> for CPUBrand {
    fn from(s: &str) -> Self {
        let sl = s.to_lowercase();
        if sl.contains("intel") {
            Self::Intel
        } else if sl.contains("amd") {
            Self::AMD
        } else if sl.contains("snapdragon") {
            Self::SnapDragon
        } else if sl.contains("apple") {
            Self::Apple
        } else {
            Self::Other(s.to_string())
        }
    }
}

impl fmt::Display for CPUBrand {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Intel => write!(f, "Intel"),
            Self::AMD => write!(f, "AMD"),
            Self::SnapDragon => write!(f, "SnapDragon"),
            Self::Apple => write!(f, "Apple"),
            Self::Other(s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CPUArch {
    X86,
    X86_64,
    AArch64,
    Other(String),
}

impl FromStr for CPUArch {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(match s {
            "x86" => Self::X86,
            "x86_64" => Self::X86_64,
            "arm" => Self::AArch64,
            "aarch64" => Self::AArch64,
            _ => Self::Other(s.to_string()),
        })
    }
}

impl fmt::Display for CPUArch {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::X86 => write!(f, "x86"),
            Self::X86_64 => write!(f, "x86_64"),
            Self::AArch64 => write!(f, "AArch64"),
            Self::Other(s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ClockSpeed(pub u32);

impl ClockSpeed {
    pub fn as_mhz(&self) -> u32 {
        self.0
    }

    pub fn as_ghz(&self) -> f32 {
        self.0 as f32 / 1000.0
    }
}

impl fmt::Display for ClockSpeed {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:.2} GHz ({} MHz)", self.as_ghz(), self.as_mhz())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CPU {
    pub brand: CPUBrand,
    pub arch: CPUArch,
    pub name: String,
    pub vendor_id: String,
    pub physical_cores: Option<u8>,
    pub logical_cores: Option<u8>,
    pub max_clock_speed: ClockSpeed,
    pub current_clock_speed: ClockSpeed,
    pub l1_cache: Option<Memory>,
    pub l2_cache: Option<Memory>,
    pub l3_cache: Option<Memory>,
}

impl CPU {
    pub fn current() -> Self {
        let mut sys = System::new_with_specifics(RefreshKind::default().with_cpu(CpuRefreshKind::everything()));
        sys.refresh_cpu_all();

        let cpu_opt = sys.cpus().first();

        let brand = cpu_opt.map(|cpu| cpu.brand().into())
            .unwrap_or(CPUBrand::Other("unknown".into()));
        let name = cpu_opt.map(|cpu| cpu.name().to_string())
            .unwrap_or_else(|| "unknown".into());
        let vendor_id = cpu_opt.map(|cpu| cpu.vendor_id().to_string())
            .unwrap_or_else(|| "unknown".into());
        let max_clock_speed = cpu_opt.map(|cpu| ClockSpeed(cpu.frequency() as u32))
            .unwrap_or(ClockSpeed(0));
        let current_clock_speed = max_clock_speed;

        let logical_cores = cpu_opt.map(|_| sys.cpus().len() as u8);
        let physical_cores = sysinfo::System::physical_core_count().map(|pc| pc as u8);

        let cpuid = CpuId::new();
        let mut l1_cache = None;
        let mut l2_cache = None;
        let mut l3_cache = None;
        if let Some(iter) = cpuid.get_cache_parameters() {
            for cache in iter {
                match cache.level() {
                    1 => {
                        let size = cache.physical_line_partitions()
                            * cache.coherency_line_size()
                            * cache.associativity();
                        if size > 0 { l1_cache = Some(Memory::from_bytes(size.try_into().unwrap())); }
                    },
                    2 => {
                        let size = cache.physical_line_partitions()
                            * cache.coherency_line_size()
                            * cache.associativity();
                        if size > 0 { l2_cache = Some(Memory::from_bytes(size.try_into().unwrap())); }
                    },
                    3 => {
                        let size = (cache.physical_line_partitions() as u64)
                            * (cache.coherency_line_size() as u64)
                            * (cache.associativity() as u64);
                        if size > 0 { l3_cache = Some(Memory::from_bytes(size)); }
                    },
                    _ => {}
                }
            }
        }

        Self {
            brand,
            arch: std::env::consts::ARCH.parse().unwrap_or(CPUArch::Other("unknown".into())),
            name,
            vendor_id,
            physical_cores,
            logical_cores,
            max_clock_speed,
            current_clock_speed,
            l1_cache,
            l2_cache,
            l3_cache,
        }
    }

    pub fn is_intel(&self) -> bool {
        matches!(self.brand, CPUBrand::Intel)
    }

    pub fn is_amd(&self) -> bool {
        matches!(self.brand, CPUBrand::AMD)
    }

    pub fn is_arm(&self) -> bool {
        matches!(self.brand,CPUBrand::Intel)
    }

    pub fn is_high_clock(&self) -> bool {
        self.max_clock_speed.0 > 3000
    }

    pub fn verbose_info(&self) -> String {
        format!(
            "CPU Information:\n\
            - Brand: {}\n\
            - Architecture: {}\n\
            - Name: {}\n\
            - Vendor ID: {}\n\
            - Cores: {} physical, {} logical\n\
            - Clock Speed: {} (max), {} (current)\n\
            - Cache: L1 {}, L2 {}, L3 {}",
            self.brand,
            self.arch,
            self.name,
            self.vendor_id,
            self.physical_cores.map(|c| c.to_string()).unwrap_or_else(|| "unknown".into()),
            self.logical_cores.map(|c| c.to_string()).unwrap_or_else(|| "unknown".into()),
            self.max_clock_speed,
            self.current_clock_speed,
            self.l1_cache.map(|c| c.format_human()).unwrap_or_else(|| "unknown".into()),
            self.l2_cache.map(|c| c.format_human()).unwrap_or_else(|| "unknown".into()),
            self.l3_cache.map(|c| c.format_human()).unwrap_or_else(|| "unknown".into()),
        )
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum GPUBrand {
    Nvidia,
    AMD,
    Intel,
    Other(String),
}

impl From<&str> for GPUBrand {
    fn from(s: &str) -> Self {
        let sl = s.to_lowercase();
        if sl.contains("nvidia") || sl.contains("geforce") {
            Self::Nvidia
        } else if sl.contains("amd") || sl.contains("radeon") {
            Self::AMD
        } else if sl.contains("intel") {
            Self::Intel
        } else {
            Self::Other(s.to_string())
        }
    }
}

impl fmt::Display for GPUBrand {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Nvidia => write!(f, "NVIDIA"),
            Self::AMD => write!(f, "AMD"),
            Self::Intel => write!(f, "Intel"),
            Self::Other(s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GPU {
    pub brand: GPUBrand,
    pub name: String,
    pub device_type: DeviceType,
    pub vram: Memory,
    pub driver_version: Option<String>,
}


impl GPU {
    pub fn current() -> Vec<Self> {
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
            backends: wgpu::Backends::all(),
            ..Default::default()
        });

        instance
            .enumerate_adapters(wgpu::Backends::all())
            .iter()
            .map(|adapter| {
                let info = adapter.get_info();
                GPU {
                    brand: info.name.as_str().into(),
                    name: info.name.to_string(),
                    device_type: info.device_type,
                    vram: Memory::from_bytes(0),
                    driver_version: Some(info.driver.to_string()),
                }
            })
            .collect()
    }

    pub fn is_integrated(&self) -> bool {
        matches!(self.brand, GPUBrand::Intel)
    }

    pub fn is_dedicated(&self) -> bool {
        !self.is_integrated()
    }
    pub fn is_mobile(&self) -> bool {
        let lower_name = self.name.to_lowercase();
        lower_name.contains("adreno")
            || lower_name.contains("mali")
            || lower_name.contains("apple")
            || lower_name.contains("mobile")
            || lower_name.contains("snapdragon")
    }

    pub fn verbose_info(&self) -> String {
        format!(
            "GPU Information:\n\
            - Brand: {}\n\
            - Name: {}\n\
            - VRAM: {}\n\
            - Driver: {}",
            self.brand,
            self.name,
            self.vram.format_human(),
            self.driver_version.as_deref().unwrap_or("Unknown")
        )
    }

    fn unique_id(&self) -> String {
        format!("{}-{:?}", self.name, self.device_type)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SystemMemory {
    pub total: Memory,
    pub used: Memory,
    pub free: Memory,
    pub available: Memory,
    pub swap_total: Memory,
    pub swap_used: Memory,
    pub swap_free: Memory,
}

impl SystemMemory {
    pub fn current() -> Self {
        let mut system = System::new();
        system.refresh_memory();

        Self {
            total: Memory::from_bytes(system.total_memory()),
            used: Memory::from_bytes(system.used_memory()),
            free: Memory::from_bytes(system.free_memory()),
            available: Memory::from_bytes(system.available_memory()),
            swap_total: Memory::from_bytes(system.total_swap()),
            swap_used: Memory::from_bytes(system.used_swap()),
            swap_free: Memory::from_bytes(system.free_swap()),
        }
    }

    pub fn verbose_info(&self) -> String {
        format!(
            "Memory Information:\n\
            - RAM: {} total, {} used, {} free, {} available\n\
            - Swap: {} total, {} used, {} free",
            self.total.format_human(),
            self.used.format_human(),
            self.free.format_human(),
            self.available.format_human(),
            self.swap_total.format_human(),
            self.swap_used.format_human(),
            self.swap_free.format_human()
        )
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompileInfo {
    pub timestamp: String,
    pub pkg_version: String,
    pub pkg_name: String,
    pub target_arch: String,
    pub target_os: String,
    pub target_env: String,
}

impl CompileInfo {
    pub fn current() -> Self {
        Self {
            timestamp: build_info::BUILT_TIME_UTC.to_string(),
            pkg_version: build_info::PKG_VERSION.to_string(),
            pkg_name: build_info::PKG_NAME.to_string(),
            target_arch: build_info::CFG_TARGET_ARCH.to_string(),
            target_os: build_info::CFG_OS.to_string(),
            target_env: build_info::CFG_ENV.to_string(),
        }
    }

    pub fn verbose_info(&self) -> String {
        format!(
            "Build Information:\n\
            - Timestamp: {}\n\
            - Package: {} v{}\n\
            - Target: {} {} {}\n\
            ",
            self.timestamp,
            self.pkg_name,
            self.pkg_version,
            self.target_arch,
            self.target_os,
            self.target_env
        )
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SystemMetadata {
    pub cpu: CPU,
    pub memory: SystemMemory,
    pub gpus: Vec<GPU>,
    pub compile_info: CompileInfo,
}

impl SystemMetadata {
    pub fn current() -> Self {
        Self {
            cpu: CPU::current(),
            memory: SystemMemory::current(),
            gpus: GPU::current(),
            compile_info: CompileInfo::current(),
        }
    }
    pub fn main_gpu(&self) -> Option<&GPU> {
        self.gpus
            .iter()
            .find(|g| g.is_dedicated())
            .or_else(|| self.gpus.iter().find(|g| g.is_integrated()))
            .or_else(|| self.gpus.first())
    }


    pub fn verbose_summary(&self) -> String {
        let main_gpu = self.main_gpu();
        let main_gpu_info = main_gpu
            .map(|gpu| format!("Main GPU:\n{}", gpu.verbose_info()))
            .unwrap_or_else(|| "No main GPU detected".to_string());

        let mut seen_gpu_ids = HashSet::new();
        let mut other_gpus_info = Vec::new();

        if let Some(gpu) = main_gpu {
            seen_gpu_ids.insert(gpu.unique_id());
        }

        for gpu in &self.gpus {
            let gpu_id = gpu.unique_id();
            if !seen_gpu_ids.contains(&gpu_id) {
                other_gpus_info.push(format!("[{:?}] {}", gpu.device_type, gpu.verbose_info()));
                seen_gpu_ids.insert(gpu_id);
            }
        }

        let other_gpu_list = if other_gpus_info.is_empty() {
            String::new()
        } else {
            format!("Other GPUs:\n{}", other_gpus_info.join("\n\n"))
        };

        format!(
            "System Metadata:\n\n{}\n\n{}\n\n{}\n\n{}\n\n{}",
            self.cpu.verbose_info(),
            main_gpu_info,
            other_gpu_list,
            self.memory.verbose_info(),
            self.compile_info.verbose_info()
        )
    }
}

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

    #[test]
    fn test_memory_conversions() {
        let mem = Memory::from_gb(8);
        assert_eq!(mem.as_bytes(), 8 * 1024 * 1024 * 1024);
        assert_eq!(mem.as_mb(), 8 * 1024);
        assert_eq!(mem.as_kb(), 8 * 1024 * 1024);
    }

    #[test]
    fn test_system_metadata() {
        let metadata = SystemMetadata::current();
        assert!(!metadata.cpu.name.is_empty());
        assert!(metadata.memory.total.as_bytes() > 0);
        assert!(!metadata.compile_info.pkg_version.is_empty());
    }
}