dirigent/crates/dirigent_inspector/src/system.rs

use crate::node::NodeId;
use crate::registry::InspectorRegistry;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use sysinfo::System;
use tokio::task::JoinHandle;
use tracing::warn;

/// Host system information.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SystemInfo {
    pub hostname: Option<String>,
    pub os_name: Option<String>,
    pub os_version: Option<String>,
    pub kernel_version: Option<String>,
    pub arch: String,
    pub total_memory_bytes: u64,
    pub used_memory_bytes: u64,
    pub available_memory_bytes: u64,
    pub total_swap_bytes: u64,
    pub used_swap_bytes: u64,
    pub cpu_count: usize,
    pub physical_core_count: Option<usize>,
    pub global_cpu_usage_percent: f32,
    pub uptime_secs: u64,
}

/// Monitors host system metrics (memory, CPU, etc.).
pub struct SystemMonitor {
    system: System,
}

impl SystemMonitor {
    /// Create a new system monitor.
    pub fn new() -> Self {
        let mut system = System::new();
        // Initial refresh to populate baseline data
        system.refresh_memory();
        system.refresh_cpu_usage();
        Self { system }
    }

    /// Refresh and return current system information.
    pub fn refresh(&mut self) -> SystemInfo {
        self.system.refresh_memory();
        self.system.refresh_cpu_usage();

        SystemInfo {
            hostname: System::host_name(),
            os_name: System::name(),
            os_version: System::os_version(),
            kernel_version: System::kernel_version(),
            arch: System::cpu_arch(),
            total_memory_bytes: self.system.total_memory(),
            used_memory_bytes: self.system.used_memory(),
            available_memory_bytes: self.system.available_memory(),
            total_swap_bytes: self.system.total_swap(),
            used_swap_bytes: self.system.used_swap(),
            cpu_count: self.system.cpus().len(),
            physical_core_count: self.system.physical_core_count(),
            global_cpu_usage_percent: self.system.global_cpu_usage(),
            uptime_secs: System::uptime(),
        }
    }

    /// Spawn a background task that periodically updates a system node in the registry.
    ///
    /// The `node_id` should already be registered in the tree (e.g., "dirigent/system/host").
    pub fn start_polling(
        mut self,
        registry: Arc<InspectorRegistry>,
        node_id: NodeId,
        interval: Duration,
    ) -> JoinHandle<()> {
        tokio::spawn(async move {
            let mut ticker = tokio::time::interval(interval);
            loop {
                ticker.tick().await;

                let info = self.refresh();

                let mut props = HashMap::new();
                props.insert("hostname".to_string(), serde_json::json!(info.hostname));
                props.insert("os_name".to_string(), serde_json::json!(info.os_name));
                props.insert("os_version".to_string(), serde_json::json!(info.os_version));
                props.insert("arch".to_string(), serde_json::json!(info.arch));
                props.insert(
                    "total_memory_bytes".to_string(),
                    serde_json::json!(info.total_memory_bytes),
                );
                props.insert(
                    "used_memory_bytes".to_string(),
                    serde_json::json!(info.used_memory_bytes),
                );
                props.insert(
                    "available_memory_bytes".to_string(),
                    serde_json::json!(info.available_memory_bytes),
                );
                props.insert(
                    "total_swap_bytes".to_string(),
                    serde_json::json!(info.total_swap_bytes),
                );
                props.insert(
                    "used_swap_bytes".to_string(),
                    serde_json::json!(info.used_swap_bytes),
                );
                props.insert("cpu_count".to_string(), serde_json::json!(info.cpu_count));
                props.insert(
                    "physical_core_count".to_string(),
                    serde_json::json!(info.physical_core_count),
                );
                props.insert(
                    "global_cpu_usage_percent".to_string(),
                    serde_json::json!(info.global_cpu_usage_percent),
                );
                props.insert(
                    "uptime_secs".to_string(),
                    serde_json::json!(info.uptime_secs),
                );

                if let Err(e) = registry.update_properties(&node_id, props).await {
                    warn!(
                        node_id = %node_id,
                        error = %e,
                        "Failed to update system node properties"
                    );
                }
            }
        })
    }
}

impl Default for SystemMonitor {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_system_monitor_refresh() {
        let mut monitor = SystemMonitor::new();

        // Need a brief sleep for CPU usage sampling
        std::thread::sleep(std::time::Duration::from_millis(200));

        let info = monitor.refresh();

        assert!(info.total_memory_bytes > 0, "Should have total memory");
        assert!(info.cpu_count > 0, "Should have at least 1 CPU");
        assert!(info.uptime_secs > 0, "Should have uptime");
        assert!(info.arch.len() > 0, "Should have arch info");
    }

    #[test]
    fn test_system_info_serialization() {
        let info = SystemInfo {
            hostname: Some("testhost".to_string()),
            os_name: Some("macOS".to_string()),
            os_version: Some("14.0".to_string()),
            kernel_version: Some("23.0.0".to_string()),
            arch: "arm64".to_string(),
            total_memory_bytes: 16 * 1024 * 1024 * 1024,
            used_memory_bytes: 8 * 1024 * 1024 * 1024,
            available_memory_bytes: 8 * 1024 * 1024 * 1024,
            total_swap_bytes: 2 * 1024 * 1024 * 1024,
            used_swap_bytes: 512 * 1024 * 1024,
            cpu_count: 10,
            physical_core_count: Some(10),
            global_cpu_usage_percent: 25.5,
            uptime_secs: 86400,
        };

        let json = serde_json::to_string(&info).unwrap();
        let deserialized: SystemInfo = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.hostname, Some("testhost".to_string()));
        assert_eq!(deserialized.total_memory_bytes, 16 * 1024 * 1024 * 1024);
        assert_eq!(deserialized.cpu_count, 10);
    }

    #[tokio::test]
    async fn test_system_monitor_polling() {
        use crate::node::{NodeKind, NodeMetadata, NodeState};

        let registry = Arc::new(InspectorRegistry::new());
        let root = registry.root_id().await;

        // Register system node
        let node_id = NodeId::new("dirigent/system/host");
        let mut handle = registry
            .register(
                node_id.clone(),
                &root,
                NodeMetadata::new(NodeKind::System, "Host System").with_state(NodeState::Running),
                None,
            )
            .await
            .unwrap();
        // Detach so polling task can update it
        // (handle would deregister on drop otherwise)
        handle.detach();

        let monitor = SystemMonitor::new();
        let task = monitor.start_polling(
            Arc::clone(&registry),
            node_id.clone(),
            Duration::from_millis(100),
        );

        // Wait for at least one poll cycle
        tokio::time::sleep(Duration::from_millis(250)).await;

        let meta = registry.get_node(&node_id).await.unwrap();
        assert!(
            meta.properties.contains_key("total_memory_bytes"),
            "Should have system metrics after polling"
        );
        assert!(
            meta.properties.contains_key("cpu_count"),
            "Should have CPU count"
        );

        task.abort();
    }
}