falsegreen-agent 0.1.2

//! Conservative host discovery used to select a compatible managed runtime.

use std::fmt;
use std::fs;
use std::path::Path;

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OperatingSystem {
    Linux,
    Macos,
    Windows,
    Unsupported,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Architecture {
    X86_64,
    Aarch64,
    Unsupported,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimePlatform {
    pub os: OperatingSystem,
    pub architecture: Architecture,
}

impl RuntimePlatform {
    #[must_use]
    pub const fn current() -> Self {
        let os = if cfg!(target_os = "linux") {
            OperatingSystem::Linux
        } else if cfg!(target_os = "macos") {
            OperatingSystem::Macos
        } else if cfg!(target_os = "windows") {
            OperatingSystem::Windows
        } else {
            OperatingSystem::Unsupported
        };
        let architecture = if cfg!(target_arch = "x86_64") {
            Architecture::X86_64
        } else if cfg!(target_arch = "aarch64") {
            Architecture::Aarch64
        } else {
            Architecture::Unsupported
        };
        Self { os, architecture }
    }

    #[must_use]
    pub const fn cache_key(self) -> &'static str {
        match (self.os, self.architecture) {
            (OperatingSystem::Linux, Architecture::X86_64) => "linux-x86_64",
            (OperatingSystem::Linux, Architecture::Aarch64) => "linux-aarch64",
            (OperatingSystem::Macos, Architecture::X86_64) => "macos-x86_64",
            (OperatingSystem::Macos, Architecture::Aarch64) => "macos-aarch64",
            (OperatingSystem::Windows, Architecture::X86_64) => "windows-x86_64",
            (OperatingSystem::Windows, Architecture::Aarch64) => "windows-aarch64",
            _ => "unsupported",
        }
    }
}

impl fmt::Display for RuntimePlatform {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(self.cache_key())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RuntimeBackend {
    Cpu,
    Metal,
    Rocm,
    Vulkan,
}

impl RuntimeBackend {
    #[must_use]
    pub const fn cache_key(self) -> &'static str {
        match self {
            Self::Cpu => "cpu",
            Self::Metal => "metal",
            Self::Rocm => "rocm",
            Self::Vulkan => "vulkan",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GpuDevice {
    pub identifier: String,
    pub vendor: String,
    #[serde(default)]
    pub vendor_id: Option<String>,
    pub device_id: Option<String>,
    #[serde(default)]
    pub subsystem_vendor_id: Option<String>,
    #[serde(default)]
    pub subsystem_device_id: Option<String>,
    #[serde(default)]
    pub architecture: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HardwareProfile {
    pub platform: RuntimePlatform,
    #[serde(default)]
    pub kernel_release: Option<String>,
    pub logical_cpus: usize,
    pub total_memory_bytes: Option<u64>,
    pub gpus: Vec<GpuDevice>,
    pub rocm_available: bool,
    pub vulkan_available: bool,
    #[serde(default)]
    pub graphics_driver: Option<String>,
    pub diagnostics: Vec<String>,
}

impl HardwareProfile {
    #[must_use]
    pub fn detect() -> Self {
        let platform = RuntimePlatform::current();
        let kernel_release = (platform.os == OperatingSystem::Linux)
            .then(|| read_trimmed(Path::new("/proc/sys/kernel/osrelease")))
            .flatten();
        let logical_cpus = std::thread::available_parallelism().map_or(1, usize::from);
        let mut diagnostics = Vec::new();
        let total_memory_bytes = detect_memory(platform, &mut diagnostics);
        let (gpus, rocm_available, vulkan_available) =
            detect_accelerators(platform, &mut diagnostics);
        let graphics_driver = detect_graphics_driver(platform, &gpus, kernel_release.as_deref());
        Self {
            platform,
            kernel_release,
            logical_cpus,
            total_memory_bytes,
            gpus,
            rocm_available,
            vulkan_available,
            graphics_driver,
            diagnostics,
        }
    }

    /// Choose the strongest backend that has both observable host support and a pinned artifact.
    /// Unknown Windows GPU state deliberately falls back to CPU rather than guessing.
    #[must_use]
    pub fn recommended_backend(&self) -> RuntimeBackend {
        if self.platform.os == OperatingSystem::Macos {
            return RuntimeBackend::Metal;
        }
        if self.platform.os == OperatingSystem::Linux && self.rocm_available {
            return RuntimeBackend::Rocm;
        }
        if self.platform.os == OperatingSystem::Linux && self.vulkan_available {
            return RuntimeBackend::Vulkan;
        }
        RuntimeBackend::Cpu
    }
}

fn detect_memory(platform: RuntimePlatform, diagnostics: &mut Vec<String>) -> Option<u64> {
    match platform.os {
        OperatingSystem::Linux => match fs::read_to_string("/proc/meminfo") {
            Ok(contents) => parse_linux_memory(&contents).or_else(|| {
                diagnostics.push("/proc/meminfo did not contain a valid MemTotal".to_owned());
                None
            }),
            Err(error) => {
                diagnostics.push(format!("could not read /proc/meminfo: {error}"));
                None
            }
        },
        OperatingSystem::Macos => macos_memory(diagnostics),
        _ => {
            diagnostics
                .push("physical memory detection is unavailable on this platform".to_owned());
            None
        }
    }
}

fn parse_linux_memory(contents: &str) -> Option<u64> {
    let line = contents
        .lines()
        .find(|line| line.starts_with("MemTotal:"))?;
    let kibibytes = line.split_whitespace().nth(1)?.parse::<u64>().ok()?;
    kibibytes.checked_mul(1024)
}

#[cfg(target_os = "macos")]
fn macos_memory(diagnostics: &mut Vec<String>) -> Option<u64> {
    let output = std::process::Command::new("/usr/sbin/sysctl")
        .args(["-n", "hw.memsize"])
        .output();
    match output {
        Ok(output) if output.status.success() => String::from_utf8(output.stdout)
            .ok()
            .and_then(|value| value.trim().parse::<u64>().ok()),
        Ok(output) => {
            diagnostics.push(format!("sysctl hw.memsize exited with {}", output.status));
            None
        }
        Err(error) => {
            diagnostics.push(format!("could not run sysctl hw.memsize: {error}"));
            None
        }
    }
}

#[cfg(not(target_os = "macos"))]
fn macos_memory(_diagnostics: &mut Vec<String>) -> Option<u64> {
    None
}

fn detect_accelerators(
    platform: RuntimePlatform,
    diagnostics: &mut Vec<String>,
) -> (Vec<GpuDevice>, bool, bool) {
    if platform.os == OperatingSystem::Macos {
        return (
            vec![GpuDevice {
                identifier: "apple-metal".to_owned(),
                vendor: "Apple".to_owned(),
                vendor_id: None,
                device_id: None,
                subsystem_vendor_id: None,
                subsystem_device_id: None,
                architecture: None,
            }],
            false,
            false,
        );
    }
    if platform.os != OperatingSystem::Linux {
        diagnostics.push(
            "GPU auto-detection is conservative on this platform; selecting CPU by default"
                .to_owned(),
        );
        return (Vec::new(), false, false);
    }

    let gpus = linux_drm_devices(Path::new("/sys/class/drm"), diagnostics);
    let rocm_available =
        Path::new("/dev/kfd").exists() && gpus.iter().any(|gpu| gpu.vendor == "AMD");
    let vulkan_available = Path::new("/dev/dri").read_dir().is_ok_and(|entries| {
        entries
            .filter_map(Result::ok)
            .any(|entry| entry.file_name().to_string_lossy().starts_with("renderD"))
    });
    (gpus, rocm_available, vulkan_available)
}

fn linux_drm_devices(root: &Path, diagnostics: &mut Vec<String>) -> Vec<GpuDevice> {
    let entries = match fs::read_dir(root) {
        Ok(entries) => entries,
        Err(error) => {
            diagnostics.push(format!("could not inspect {}: {error}", root.display()));
            return Vec::new();
        }
    };
    let mut devices = Vec::new();
    for entry in entries.filter_map(Result::ok) {
        let name = entry.file_name().to_string_lossy().into_owned();
        if !name.starts_with("card") || name.contains('-') {
            continue;
        }
        let vendor_path = entry.path().join("device/vendor");
        let Ok(vendor_id) = fs::read_to_string(&vendor_path) else {
            continue;
        };
        let vendor_id = vendor_id.trim().to_ascii_lowercase();
        let vendor = match vendor_id.as_str() {
            "0x1002" => "AMD",
            "0x10de" => "NVIDIA",
            "0x8086" => "Intel",
            _ => vendor_id.as_str(),
        };
        let device_id = fs::read_to_string(entry.path().join("device/device"))
            .ok()
            .map(|value| value.trim().to_ascii_lowercase());
        let subsystem_vendor_id = read_trimmed(&entry.path().join("device/subsystem_vendor"))
            .map(|value| value.to_ascii_lowercase());
        let subsystem_device_id = read_trimmed(&entry.path().join("device/subsystem_device"))
            .map(|value| value.to_ascii_lowercase());
        let architecture = known_gpu_architecture(&vendor_id, device_id.as_deref());
        devices.push(GpuDevice {
            identifier: name,
            vendor: vendor.to_owned(),
            vendor_id: Some(vendor_id),
            device_id,
            subsystem_vendor_id,
            subsystem_device_id,
            architecture,
        });
    }
    devices.sort_by(|left, right| left.identifier.cmp(&right.identifier));
    devices
}

fn read_trimmed(path: &Path) -> Option<String> {
    fs::read_to_string(path)
        .ok()
        .map(|value| value.trim().to_owned())
        .filter(|value| !value.is_empty())
}

fn known_gpu_architecture(vendor_id: &str, device_id: Option<&str>) -> Option<String> {
    match (vendor_id, device_id) {
        ("0x1002", Some("0x744c")) => Some("gfx1100".to_owned()),
        _ => None,
    }
}

fn detect_graphics_driver(
    platform: RuntimePlatform,
    gpus: &[GpuDevice],
    kernel_release: Option<&str>,
) -> Option<String> {
    if platform.os == OperatingSystem::Macos {
        return Some("Metal".to_owned());
    }
    if platform.os != OperatingSystem::Linux {
        return None;
    }
    let driver = if gpus.iter().any(|gpu| gpu.vendor == "AMD")
        && Path::new("/sys/module/amdgpu").exists()
    {
        "amdgpu"
    } else if gpus.iter().any(|gpu| gpu.vendor == "NVIDIA")
        && Path::new("/sys/module/nvidia").exists()
    {
        "nvidia"
    } else if gpus.iter().any(|gpu| gpu.vendor == "Intel") && Path::new("/sys/module/i915").exists()
    {
        "i915"
    } else {
        return None;
    };
    Some(kernel_release.map_or_else(
        || driver.to_owned(),
        |release| format!("{driver} on kernel {release}"),
    ))
}

#[cfg(test)]
mod tests {
    use super::{
        Architecture, GpuDevice, HardwareProfile, OperatingSystem, RuntimeBackend, RuntimePlatform,
        parse_linux_memory,
    };

    #[test]
    fn parses_linux_memtotal_as_bytes() {
        assert_eq!(
            parse_linux_memory("MemFree: 1 kB\nMemTotal:       16384 kB\n"),
            Some(16 * 1024 * 1024)
        );
        assert_eq!(parse_linux_memory("MemFree: 1 kB\n"), None);
    }

    #[test]
    fn backend_selection_requires_observable_support() {
        let base = HardwareProfile {
            platform: RuntimePlatform {
                os: OperatingSystem::Linux,
                architecture: Architecture::X86_64,
            },
            kernel_release: Some("test-kernel".to_owned()),
            logical_cpus: 8,
            total_memory_bytes: Some(32 * 1024 * 1024 * 1024),
            gpus: vec![GpuDevice {
                identifier: "card0".to_owned(),
                vendor: "AMD".to_owned(),
                vendor_id: Some("0x1002".to_owned()),
                device_id: Some("0x744c".to_owned()),
                subsystem_vendor_id: Some("0x1eae".to_owned()),
                subsystem_device_id: Some("0x7901".to_owned()),
                architecture: Some("gfx1100".to_owned()),
            }],
            rocm_available: false,
            vulkan_available: false,
            graphics_driver: Some("amdgpu on kernel test-kernel".to_owned()),
            diagnostics: Vec::new(),
        };
        assert_eq!(base.recommended_backend(), RuntimeBackend::Cpu);
        assert_eq!(
            HardwareProfile {
                vulkan_available: true,
                ..base.clone()
            }
            .recommended_backend(),
            RuntimeBackend::Vulkan
        );
        assert_eq!(
            HardwareProfile {
                rocm_available: true,
                vulkan_available: true,
                ..base
            }
            .recommended_backend(),
            RuntimeBackend::Rocm
        );
    }

    #[test]
    fn apple_platform_selects_metal() {
        let profile = HardwareProfile {
            platform: RuntimePlatform {
                os: OperatingSystem::Macos,
                architecture: Architecture::Aarch64,
            },
            kernel_release: None,
            logical_cpus: 8,
            total_memory_bytes: None,
            gpus: Vec::new(),
            rocm_available: false,
            vulkan_available: false,
            graphics_driver: Some("Metal".to_owned()),
            diagnostics: Vec::new(),
        };
        assert_eq!(profile.recommended_backend(), RuntimeBackend::Metal);
    }
}