all-smi 0.26.2

// Copyright 2025 Lablup Inc. and Jeongkyu Shin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Chassis metrics exporter for Prometheus
//!
//! Exports node-level metrics including:
//! - Total power consumption (CPU+GPU+ANE on Apple Silicon, SMC PSTR on Intel Macs)
//! - Thermal pressure (macOS)
//! - Individual power components (CPU, GPU, ANE), whichever the platform reports
//! - Fan speeds, where the platform reports them

use super::{MetricBuilder, MetricExporter};
use crate::device::ChassisInfo;

/// Exporter for chassis-level metrics
pub struct ChassisMetricExporter<'a> {
    chassis_info: &'a [ChassisInfo],
}

impl<'a> ChassisMetricExporter<'a> {
    pub fn new(chassis_info: &'a [ChassisInfo]) -> Self {
        Self { chassis_info }
    }
}

/// Flags for which metric types are present across chassis info
struct MetricPresenceFlags {
    has_power: bool,
    has_thermal_pressure: bool,
    has_cpu_power: bool,
    has_gpu_power: bool,
    has_ane_power: bool,
    has_inlet_temp: bool,
    has_outlet_temp: bool,
    has_fan_speeds: bool,
}

impl MetricPresenceFlags {
    /// Scan chassis info once to determine which metrics are present
    fn from_chassis_info(chassis_info: &[ChassisInfo]) -> Self {
        let mut flags = Self {
            has_power: false,
            has_thermal_pressure: false,
            has_cpu_power: false,
            has_gpu_power: false,
            has_ane_power: false,
            has_inlet_temp: false,
            has_outlet_temp: false,
            has_fan_speeds: false,
        };

        for chassis in chassis_info {
            flags.has_power |= chassis.total_power_watts.is_some();
            flags.has_thermal_pressure |= chassis.thermal_pressure.is_some();
            flags.has_cpu_power |= chassis.detail.contains_key("cpu_power_watts");
            flags.has_gpu_power |= chassis.detail.contains_key("gpu_power_watts");
            flags.has_ane_power |= chassis.detail.contains_key("ane_power_watts");
            flags.has_inlet_temp |= chassis.inlet_temperature.is_some();
            flags.has_outlet_temp |= chassis.outlet_temperature.is_some();
            flags.has_fan_speeds |= !chassis.fan_speeds.is_empty();

            // Early exit if all flags are set
            if flags.all_present() {
                break;
            }
        }

        flags
    }

    fn all_present(&self) -> bool {
        self.has_power
            && self.has_thermal_pressure
            && self.has_cpu_power
            && self.has_gpu_power
            && self.has_ane_power
            && self.has_inlet_temp
            && self.has_outlet_temp
            && self.has_fan_speeds
    }
}

impl<'a> MetricExporter for ChassisMetricExporter<'a> {
    fn export_metrics(&self) -> String {
        let mut builder = MetricBuilder::new();

        if self.chassis_info.is_empty() {
            return builder.build();
        }

        // Single pass to determine which metrics are present
        let flags = MetricPresenceFlags::from_chassis_info(self.chassis_info);

        // Export chassis info metric with DMI/platform details as labels
        {
            // Known detail keys to promote to Prometheus labels (display_key, label_name)
            let detail_keys: &[(&str, &str)] = &[
                ("Product Name", "product_name"),
                ("Vendor", "vendor"),
                ("Board", "board"),
                ("Version", "version"),
                ("BIOS Version", "bios_version"),
                ("platform", "platform"),
            ];
            let has_details = self
                .chassis_info
                .iter()
                .any(|c| detail_keys.iter().any(|(k, _)| c.detail.contains_key(*k)));

            if has_details {
                builder
                    .help(
                        "all_smi_chassis_info",
                        "Chassis/node identification information",
                    )
                    .type_("all_smi_chassis_info", "gauge");

                for chassis in self.chassis_info {
                    // Collect label values that are present
                    let mut label_values: Vec<(&str, &str)> = vec![
                        ("hostname", &chassis.hostname),
                        ("instance", &chassis.instance),
                    ];
                    for &(display_key, label_name) in detail_keys {
                        if let Some(val) = chassis.detail.get(display_key) {
                            label_values.push((label_name, val));
                        }
                    }
                    builder.metric("all_smi_chassis_info", &label_values, "1");
                }
            }
        }

        // Export chassis power metrics
        if flags.has_power {
            builder
                .help(
                    "all_smi_chassis_power_watts",
                    "Total chassis power consumption in watts",
                )
                .type_("all_smi_chassis_power_watts", "gauge");

            for chassis in self.chassis_info {
                if let Some(power) = chassis.total_power_watts {
                    builder.metric(
                        "all_smi_chassis_power_watts",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{power:.2}"),
                    );
                }
            }
        }

        // Export thermal pressure metric (macOS)
        if flags.has_thermal_pressure {
            builder
                .help(
                    "all_smi_chassis_thermal_pressure_info",
                    "Thermal pressure level (macOS)",
                )
                .type_("all_smi_chassis_thermal_pressure_info", "gauge");

            for chassis in self.chassis_info {
                if let Some(ref pressure) = chassis.thermal_pressure {
                    builder.metric(
                        "all_smi_chassis_thermal_pressure_info",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                            ("level", pressure),
                        ],
                        "1",
                    );
                }
            }
        }

        // Export individual power components if available
        if flags.has_cpu_power {
            builder
                .help(
                    "all_smi_chassis_cpu_power_watts",
                    "CPU power consumption in watts",
                )
                .type_("all_smi_chassis_cpu_power_watts", "gauge");

            for chassis in self.chassis_info {
                if let Some(power_str) = chassis.detail.get("cpu_power_watts")
                    && let Ok(power) = power_str.parse::<f64>()
                {
                    builder.metric(
                        "all_smi_chassis_cpu_power_watts",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{power:.2}"),
                    );
                }
            }
        }

        if flags.has_gpu_power {
            builder
                .help(
                    "all_smi_chassis_gpu_power_watts",
                    "GPU power consumption in watts",
                )
                .type_("all_smi_chassis_gpu_power_watts", "gauge");

            for chassis in self.chassis_info {
                if let Some(power_str) = chassis.detail.get("gpu_power_watts")
                    && let Ok(power) = power_str.parse::<f64>()
                {
                    builder.metric(
                        "all_smi_chassis_gpu_power_watts",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{power:.2}"),
                    );
                }
            }
        }

        if flags.has_ane_power {
            builder
                .help(
                    "all_smi_chassis_ane_power_watts",
                    "ANE (Apple Neural Engine) power consumption in watts",
                )
                .type_("all_smi_chassis_ane_power_watts", "gauge");

            for chassis in self.chassis_info {
                if let Some(power_str) = chassis.detail.get("ane_power_watts")
                    && let Ok(power) = power_str.parse::<f64>()
                {
                    builder.metric(
                        "all_smi_chassis_ane_power_watts",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{power:.2}"),
                    );
                }
            }
        }

        // Export inlet/outlet temperature if available
        if flags.has_inlet_temp {
            builder
                .help(
                    "all_smi_chassis_inlet_temperature_celsius",
                    "Chassis inlet temperature in Celsius",
                )
                .type_("all_smi_chassis_inlet_temperature_celsius", "gauge");

            for chassis in self.chassis_info {
                if let Some(temp) = chassis.inlet_temperature {
                    builder.metric(
                        "all_smi_chassis_inlet_temperature_celsius",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{temp:.1}"),
                    );
                }
            }
        }

        if flags.has_outlet_temp {
            builder
                .help(
                    "all_smi_chassis_outlet_temperature_celsius",
                    "Chassis outlet temperature in Celsius",
                )
                .type_("all_smi_chassis_outlet_temperature_celsius", "gauge");

            for chassis in self.chassis_info {
                if let Some(temp) = chassis.outlet_temperature {
                    builder.metric(
                        "all_smi_chassis_outlet_temperature_celsius",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                        ],
                        format!("{temp:.1}"),
                    );
                }
            }
        }

        // Export fan speed metrics if available
        if flags.has_fan_speeds {
            builder
                .help("all_smi_chassis_fan_speed_rpm", "Fan speed in RPM")
                .type_("all_smi_chassis_fan_speed_rpm", "gauge");

            for chassis in self.chassis_info {
                for fan in &chassis.fan_speeds {
                    builder.metric(
                        "all_smi_chassis_fan_speed_rpm",
                        &[
                            ("hostname", &chassis.hostname),
                            ("instance", &chassis.instance),
                            ("fan_id", &fan.id.to_string()),
                            ("fan_name", &fan.name),
                        ],
                        fan.speed_rpm.to_string(),
                    );
                }
            }
        }

        builder.build()
    }
}

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

    #[test]
    fn test_empty_chassis_info() {
        let exporter = ChassisMetricExporter::new(&[]);
        let metrics = exporter.export_metrics();
        assert!(metrics.is_empty());
    }

    #[test]
    fn test_chassis_power_metric() {
        let chassis = ChassisInfo {
            hostname: "test-host".to_string(),
            instance: "test-instance".to_string(),
            total_power_watts: Some(45.5),
            ..Default::default()
        };

        let chassis_vec = vec![chassis];
        let exporter = ChassisMetricExporter::new(&chassis_vec);
        let metrics = exporter.export_metrics();

        assert!(metrics.contains("all_smi_chassis_power_watts"));
        assert!(metrics.contains("hostname=\"test-host\""));
        assert!(metrics.contains("45.50"));
    }

    #[test]
    fn test_thermal_pressure_metric() {
        let chassis = ChassisInfo {
            hostname: "mac-host".to_string(),
            instance: "mac-instance".to_string(),
            thermal_pressure: Some("Nominal".to_string()),
            ..Default::default()
        };

        let chassis_vec = vec![chassis];
        let exporter = ChassisMetricExporter::new(&chassis_vec);
        let metrics = exporter.export_metrics();

        assert!(metrics.contains("all_smi_chassis_thermal_pressure_info"));
        assert!(metrics.contains("level=\"Nominal\""));
    }

    #[test]
    fn test_chassis_info_dmi_labels_metric() {
        let mut detail = std::collections::HashMap::new();
        detail.insert("Product Name".to_string(), "DGX H100".to_string());
        detail.insert("Vendor".to_string(), "NVIDIA".to_string());
        detail.insert("Board".to_string(), "H100-BOARD".to_string());
        detail.insert("BIOS Version".to_string(), "1.0.0".to_string());
        detail.insert("platform".to_string(), "Linux".to_string());

        let chassis = ChassisInfo {
            hostname: "dgx-host".to_string(),
            instance: "dgx-instance".to_string(),
            detail,
            ..Default::default()
        };

        let chassis_vec = vec![chassis];
        let exporter = ChassisMetricExporter::new(&chassis_vec);
        let metrics = exporter.export_metrics();

        assert!(metrics.contains("all_smi_chassis_info"));
        assert!(metrics.contains("product_name=\"DGX H100\""));
        assert!(metrics.contains("vendor=\"NVIDIA\""));
        assert!(metrics.contains("board=\"H100-BOARD\""));
        assert!(metrics.contains("bios_version=\"1.0.0\""));
        assert!(metrics.contains("platform=\"Linux\""));
        assert!(metrics.contains("hostname=\"dgx-host\""));
        // The metric value should be 1
        assert!(metrics.contains("} 1\n"));
    }

    #[test]
    fn test_chassis_inlet_outlet_temperature_metrics() {
        let chassis = ChassisInfo {
            hostname: "server1".to_string(),
            instance: "server1".to_string(),
            inlet_temperature: Some(22.5),
            outlet_temperature: Some(35.0),
            ..Default::default()
        };

        let chassis_vec = vec![chassis];
        let exporter = ChassisMetricExporter::new(&chassis_vec);
        let metrics = exporter.export_metrics();

        assert!(metrics.contains("all_smi_chassis_inlet_temperature_celsius"));
        assert!(metrics.contains("22.5"));
        assert!(metrics.contains("all_smi_chassis_outlet_temperature_celsius"));
        assert!(metrics.contains("35.0"));
    }

    #[test]
    fn test_chassis_info_no_dmi_details_skips_info_metric() {
        // A chassis with no recognized detail keys should not emit all_smi_chassis_info
        let chassis = ChassisInfo {
            hostname: "bare-host".to_string(),
            instance: "bare-instance".to_string(),
            ..Default::default()
        };

        let chassis_vec = vec![chassis];
        let exporter = ChassisMetricExporter::new(&chassis_vec);
        let metrics = exporter.export_metrics();

        assert!(!metrics.contains("all_smi_chassis_info"));
    }

    #[test]
    fn test_metric_presence_flags_all_present() {
        let mut detail = std::collections::HashMap::new();
        detail.insert("cpu_power_watts".to_string(), "15.0".to_string());
        detail.insert("gpu_power_watts".to_string(), "200.0".to_string());
        detail.insert("ane_power_watts".to_string(), "5.0".to_string());

        let chassis = ChassisInfo {
            hostname: "full-host".to_string(),
            instance: "full-instance".to_string(),
            total_power_watts: Some(220.0),
            thermal_pressure: Some("Nominal".to_string()),
            inlet_temperature: Some(20.0),
            outlet_temperature: Some(30.0),
            fan_speeds: vec![crate::device::FanInfo {
                id: 0,
                name: "Fan0".to_string(),
                speed_rpm: 1200,
                max_rpm: 3000,
            }],
            detail,
            ..Default::default()
        };

        let flags = MetricPresenceFlags::from_chassis_info(&[chassis]);
        assert!(flags.all_present());
    }
}