ember-tune-rs/src/sal/heuristic/discovery.rs

use std::fs;
use std::path::{Path, PathBuf};
use std::time::{Duration};
use std::thread;
use std::sync::mpsc;
use std::collections::HashMap;
use crate::sal::heuristic::schema::{SensorDiscovery, ActuatorDiscovery, Conflict, Discovery, Benchmarking};
use crate::sys::SyscallRunner;
use tracing::{debug, warn};

/// Registry of dynamically discovered paths for configs and tools.
#[derive(Debug, Clone, Default)]
pub struct PathRegistry {
    pub configs: HashMap<String, PathBuf>,
    pub tools: HashMap<String, PathBuf>,
}

/// Strongly-typed findings about the current system.
#[derive(Debug, Clone, Default)]
pub struct SystemFactSheet {
    pub vendor: String,
    pub model: String,
    pub temp_path: Option<PathBuf>,
    pub fan_paths: Vec<PathBuf>,
    pub rapl_paths: Vec<PathBuf>,
    pub active_conflicts: Vec<String>,
    pub paths: PathRegistry,
    pub bench_config: Option<Benchmarking>,
}

/// Probes the system for hardware sensors, actuators, service conflicts, and paths.
pub fn discover_facts(
    base_path: &Path,
    runner: &dyn SyscallRunner,
    discovery: &Discovery,
    conflicts: &[Conflict],
    bench_config: Benchmarking,
) -> SystemFactSheet {
    let (vendor, model) = read_dmi_info(base_path);
    
    debug!("DMI Identity: Vendor='{}', Model='{}'", vendor, model);

    let (temp_path, fan_paths) = discover_hwmon(base_path, &discovery.sensors);
    let rapl_paths = discover_rapl(base_path, &discovery.actuators);
    
    let mut active_conflicts = Vec::new();
    for conflict in conflicts {
        for service in &conflict.services {
            if is_service_active(runner, service) {
                debug!("Detected active conflict: {} (Service: {})", conflict.id, service);
                active_conflicts.push(conflict.id.clone());
                break;
            }
        }
    }

    let paths = discover_paths(base_path, discovery);

    SystemFactSheet {
        vendor,
        model,
        temp_path,
        fan_paths,
        rapl_paths,
        active_conflicts,
        paths,
        bench_config: Some(bench_config),
    }
}

fn discover_paths(base_path: &Path, discovery: &Discovery) -> PathRegistry {
    let mut registry = PathRegistry::default();

    // 1. Discover Tools via PATH
    for (id, binary_name) in &discovery.tools {
        if let Ok(path) = which::which(binary_name) {
            debug!("Discovered tool: {} -> {:?}", id, path);
            registry.tools.insert(id.clone(), path);
        }
    }

    // 2. Discover Configs via existence check
    for (id, candidates) in &discovery.configs {
        for candidate in candidates {
            let path = if candidate.starts_with('/') {
                base_path.join(&candidate[1..])
            } else {
                base_path.join(candidate)
            };
            
            if path.exists() {
                debug!("Discovered config: {} -> {:?}", id, path);
                registry.configs.insert(id.clone(), path);
                break;
            }
        }
        if !registry.configs.contains_key(id) {
            if let Some(first) = candidates.first() {
                registry.configs.insert(id.clone(), PathBuf::from(first));
            }
        }
    }

    registry
}

/// Reads DMI information from sysfs with a safety timeout.
fn read_dmi_info(base_path: &Path) -> (String, String) {
    let vendor = read_sysfs_with_timeout(&base_path.join("sys/class/dmi/id/sys_vendor"), Duration::from_millis(100))
        .unwrap_or_else(|| "Unknown".to_string());
    let model = read_sysfs_with_timeout(&base_path.join("sys/class/dmi/id/product_name"), Duration::from_millis(100))
        .unwrap_or_else(|| "Unknown".to_string());
    (vendor, model)
}

/// Discovers hwmon sensors by matching labels and prioritizing drivers.
fn discover_hwmon(base_path: &Path, cfg: &SensorDiscovery) -> (Option<PathBuf>, Vec<PathBuf>) {
    let mut temp_candidates = Vec::new();
    let mut fan_candidates = Vec::new();

    let hwmon_base = base_path.join("sys/class/hwmon");
    let entries = match fs::read_dir(&hwmon_base) {
        Ok(e) => e,
        Err(e) => {
            warn!("Could not read {:?}: {}", hwmon_base, e);
            return (None, Vec::new());
        }
    };

    for entry in entries.flatten() {
        let hwmon_path = entry.path();
        
        let driver_name = read_sysfs_with_timeout(&hwmon_path.join("name"), Duration::from_millis(100))
            .unwrap_or_default();

        let priority = cfg.hwmon_priority
            .iter()
            .position(|p| p == &driver_name)
            .unwrap_or(usize::MAX);

        if let Ok(hw_entries) = fs::read_dir(&hwmon_path) {
            for hw_entry in hw_entries.flatten() {
                let file_name = hw_entry.file_name().into_string().unwrap_or_default();
                
                if file_name.starts_with("temp") && file_name.ends_with("_label") {
                    if let Some(label) = read_sysfs_with_timeout(&hw_entry.path(), Duration::from_millis(100)) {
                        if cfg.temp_labels.iter().any(|l| label.contains(l)) {
                            let input_path = hwmon_path.join(file_name.replace("_label", "_input"));
                            if input_path.exists() {
                                temp_candidates.push((priority, input_path));
                            }
                        }
                    }
                }

                if file_name.starts_with("fan") && file_name.ends_with("_label") {
                    if let Some(label) = read_sysfs_with_timeout(&hw_entry.path(), Duration::from_millis(100)) {
                        if cfg.fan_labels.iter().any(|l| label.contains(l)) {
                            let input_path = hwmon_path.join(file_name.replace("_label", "_input"));
                            if input_path.exists() {
                                fan_candidates.push((priority, input_path));
                            }
                        }
                    }
                }
            }
        }
    }

    temp_candidates.sort_by_key(|(p, _)| *p);
    fan_candidates.sort_by_key(|(p, _)| *p);

    let best_temp = temp_candidates.first().map(|(_, p)| p.clone());
    let best_fans = fan_candidates.into_iter().map(|(_, p)| p).collect();

    (best_temp, best_fans)
}

/// Discovers RAPL powercap paths.
fn discover_rapl(base_path: &Path, cfg: &ActuatorDiscovery) -> Vec<PathBuf> {
    let mut paths = Vec::new();
    let powercap_base = base_path.join("sys/class/powercap");
    
    let entries = match fs::read_dir(&powercap_base) {
        Ok(e) => e,
        Err(_) => return Vec::new(),
    };

    for entry in entries.flatten() {
        let path = entry.path();
        let dir_name = entry.file_name().into_string().unwrap_or_default();
        
        if cfg.rapl_paths.contains(&dir_name) {
            paths.push(path);
            continue;
        }

        if let Some(name) = read_sysfs_with_timeout(&path.join("name"), Duration::from_millis(100)) {
            if cfg.rapl_paths.iter().any(|p| p == &name) {
                paths.push(path);
            }
        }
    }
    paths
}

/// Checks if a systemd service is currently active using the injected runner.
pub fn is_service_active(runner: &dyn SyscallRunner, service: &str) -> bool {
    runner.run("systemctl", &["is-active", "--quiet", service]).is_ok()
}

/// Helper to read a sysfs file with a timeout.
fn read_sysfs_with_timeout(path: &Path, timeout: Duration) -> Option<String> {
    let (tx, rx) = mpsc::channel();
    let path_buf = path.to_path_buf();
    
    thread::spawn(move || {
        let res = fs::read_to_string(path_buf).map(|s| s.trim().to_string());
        let _ = tx.send(res);
    });

    match rx.recv_timeout(timeout) {
        Ok(Ok(content)) => Some(content),
        _ => None,
    }
}