Lumotia/crates/core/src/hardware.rs

use sysinfo::System;

use crate::types::Megabytes;

/// Detected system capabilities.
#[derive(Debug, Clone)]
pub struct SystemProfile {
    pub ram: Megabytes,
    pub cpu: CpuInfo,
    pub gpu: Option<GpuInfo>,
    pub os: Os,
}

#[derive(Debug, Clone)]
pub struct CpuInfo {
    pub logical_processors: usize,
    pub brand: String,
    pub features: CpuFeatures,
}

/// Runtime-detected CPU feature flags relevant to the speech-to-text
/// and LLM backends Magnotia ships. All whisper.cpp / llama.cpp / ggml
/// kernels degrade roughly two tiers without AVX2, which is why we
/// surface it separately: when AVX2 is absent, the UI should warn the
/// user that performance will be a fraction of what they would see
/// on a contemporary CPU. References:
///   - whisper-rs #8, #117 (illegal instruction on pre-AVX2 CPUs)
///   - Buzz FAQ (non-AVX2 fallback builds)
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CpuFeatures {
    pub avx2: bool,
    pub avx512f: bool,
    pub fma: bool,
    pub sse4_2: bool,
    pub neon: bool,
}

impl CpuFeatures {
    /// Whether this CPU has the baseline ggml expects (AVX2 + FMA on
    /// x86_64, NEON on aarch64). If false, the runtime banner fires.
    pub fn has_ggml_baseline(&self) -> bool {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        {
            return self.avx2 && self.fma;
        }
        #[cfg(target_arch = "aarch64")]
        {
            return self.neon;
        }
        #[allow(unreachable_code)]
        false
    }
}

/// Probes CPU feature flags via compile-time/runtime CPUID. On x86_64
/// we rely on `std::is_x86_feature_detected!`, which lowers to CPUID
/// at runtime. On aarch64 we assume NEON (architectural baseline);
/// on other targets all flags are false.
pub fn probe_cpu_features() -> CpuFeatures {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        return CpuFeatures {
            avx2: std::is_x86_feature_detected!("avx2"),
            avx512f: std::is_x86_feature_detected!("avx512f"),
            fma: std::is_x86_feature_detected!("fma"),
            sse4_2: std::is_x86_feature_detected!("sse4.2"),
            neon: false,
        };
    }
    #[cfg(target_arch = "aarch64")]
    {
        return CpuFeatures {
            avx2: false,
            avx512f: false,
            fma: false,
            sse4_2: false,
            neon: true,
        };
    }
    #[allow(unreachable_code)]
    CpuFeatures::default()
}

#[derive(Debug, Clone)]
pub struct GpuInfo {
    pub vendor: GpuVendor,
    pub vram: Megabytes,
    pub acceleration: GpuAcceleration,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GpuVendor {
    Nvidia,
    Amd,
    Intel,
    Apple,
    Unknown,
}

#[derive(Debug, Clone)]
pub struct GpuAcceleration {
    pub cuda: bool,
    pub metal: bool,
    pub vulkan: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Os {
    Windows,
    Linux,
    MacOs,
    Ios,
    Android,
}

/// Probes RAM from a shared `System` instance.
fn probe_ram_from(sys: &System) -> Megabytes {
    let total_bytes = sys.total_memory();
    Megabytes(total_bytes / (1024 * 1024))
}

/// Probes CPU info from a shared `System` instance.
fn probe_cpu_from(sys: &System) -> CpuInfo {
    CpuInfo {
        logical_processors: sys.cpus().len(),
        brand: sys
            .cpus()
            .first()
            .map(|c| c.brand().to_string())
            .unwrap_or_default(),
        features: probe_cpu_features(),
    }
}

pub fn probe_gpu() -> Option<GpuInfo> {
    // GPU detection via wgpu or platform-specific APIs.
    // Placeholder: returns None until wgpu or nvml integration is added.
    None
}

pub fn probe_os() -> Os {
    #[cfg(target_os = "windows")]
    return Os::Windows;
    #[cfg(target_os = "linux")]
    return Os::Linux;
    #[cfg(target_os = "macos")]
    return Os::MacOs;
    #[cfg(target_os = "ios")]
    return Os::Ios;
    #[cfg(target_os = "android")]
    return Os::Android;

    // Fallback for unsupported targets — treat as Linux since
    // most exotic/embedded targets are Unix-like.
    #[allow(unreachable_code)]
    Os::Linux
}

/// Composes the individual probes using a single `System` snapshot.
/// `System::new_all()` is expensive — calling it once rather than
/// per-probe avoids redundant OS queries.
pub fn probe_system() -> SystemProfile {
    let sys = System::new_all();
    SystemProfile {
        ram: probe_ram_from(&sys),
        cpu: probe_cpu_from(&sys),
        gpu: probe_gpu(),
        os: probe_os(),
    }
}

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

    #[test]
    fn probe_cpu_features_runs_without_panicking() {
        let _ = probe_cpu_features();
    }

    #[test]
    fn probe_system_populates_cpu_features() {
        let profile = probe_system();
        // The check doesn't assume the runner has AVX2; it just asserts
        // that the feature probe was actually called and is wired in.
        let f = profile.cpu.features;
        assert!(
            f == f,
            "CpuFeatures must be PartialEq so the runtime banner can debounce"
        );
    }

    #[test]
    fn ggml_baseline_matches_x86_64_rule() {
        let features = CpuFeatures {
            avx2: true,
            fma: true,
            ..CpuFeatures::default()
        };
        // Only actually true on x86_64 — on other arches the helper
        // returns false, which is equally fine for this test.
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        assert!(features.has_ggml_baseline());
        #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
        assert!(!features.has_ggml_baseline());
    }

    #[test]
    fn ggml_baseline_requires_both_avx2_and_fma() {
        let features = CpuFeatures {
            avx2: true,
            fma: false,
            ..CpuFeatures::default()
        };
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        assert!(!features.has_ggml_baseline());
        #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
        assert!(!features.has_ggml_baseline());
    }
}