Hassliebe/src/modules/info.rs

use std::fs;
use std::time::Duration;

use anyhow::{anyhow, Result};
use rumqttc::QoS;
use serde::{Deserialize, Serialize};
use sysinfo::{CpuExt, System, SystemExt};
use tokio::time::MissedTickBehavior;
use validator::Validate;

use crate::modules::InitializationContext;

const MODULE_ID: &str = "info";

#[derive(Serialize, Deserialize, Validate)]
pub struct Config {
    #[serde(default)]
    enabled: bool,

    #[serde(default)]
    ram_usage: u64,

    #[serde(default)]
    cpu_usage: u64,

    #[serde(default)]
    battery: u64,
}

pub async fn init(context: &mut InitializationContext) -> Result<()> {
    let full_context = context.get_full();
    let config = match &full_context.config.modules.info {
        Some(c) if c.enabled => c,
        _ => return Ok(()),
    };

    log::info!("Initializing…");

    {
        let mut sys = System::new();
        init_info_value(
            InfoEntityOptions {
                sub_id: "ram_usage",
                display_name: "RAM Usage",
                icon: "mdi:memory",
                suggested_precision: 0,
                unit: Some("%"),
            },
            context,
            config.ram_usage,
            move || {
                sys.refresh_memory();
                Ok(format!("{:.2}", (sys.available_memory() as f64) / (sys.total_memory() as f64) * 100.0))
            },
        )
            .await?;
    }

    {
        let mut sys = System::new();
        init_info_value(
            InfoEntityOptions {
                sub_id: "cpu_usage",
                display_name: "CPU Usage",
                icon: "mdi:memory",
                suggested_precision: 0,
                unit: Some("%"),
            },
            context,
            config.cpu_usage,
            move || {
                sys.refresh_cpu();
                Ok(format!("{:.2}", sys.global_cpu_info().cpu_usage()))
            },
        )
            .await?;
    }

    if config.battery != 0 {
        let battery_dirs = fs::read_dir("/sys/class/power_supply")?.filter_map(|d| d.ok()).collect::<Vec<_>>();

        if let Some(dir) = battery_dirs.first() {
            log::debug!("Found {} batteries, using {}", battery_dirs.len(), dir.file_name().to_string_lossy());
            let path = dir.path();
            let capacity_path = path.clone().join("capacity");
            let status_path = path.clone().join("status");

            init_info_value(
                InfoEntityOptions {
                    sub_id: "battery_level",
                    display_name: "Battery Level",
                    icon: "mdi:battery",
                    suggested_precision: 0,
                    unit: Some("%"),
                },
                context,
                config.battery,
                move || Ok(fs::read_to_string(&capacity_path)?),
            )
                .await?;

            init_info_value(
                InfoEntityOptions {
                    sub_id: "battery_state",
                    display_name: "Battery State",
                    icon: "mdi:battery",
                    suggested_precision: 0,
                    unit: None,
                },
                context,
                config.battery,
                move || Ok(fs::read_to_string(&status_path)?),
            )
                .await?;
        } else {
            log::warn!("No batteries found.")
        }
    }

    Ok(())
}

struct InfoEntityOptions<'a> {
    sub_id: &'a str,
    display_name: &'a str,
    icon: &'a str,
    suggested_precision: u64,
    unit: Option<&'a str>,
}

async fn init_info_value(
    options: InfoEntityOptions<'_>,
    context: &mut InitializationContext,
    interval_secs: u64,
    mut fetch: impl (FnMut() -> Result<String>) + Send + 'static,
) -> Result<()> {
    let interval = if interval_secs == 0 {
        return Ok(());
    } else {
        Duration::from_secs(interval_secs)
    };

    let entity_id = context.full.get_entity_id(MODULE_ID, options.sub_id);
    let state_topic = context.full.mqtt.get_homeassistant_topic("sensor", &entity_id, "state");
    context.register_owned_mqtt_topic(&state_topic);

    context
        .publish_to_owned_mqtt_topic(
            &context.full.mqtt.get_homeassistant_topic("sensor", &entity_id, "config"),
            json::stringify({
                let mut object = json::object! {
                    "availability_topic": context.full.mqtt.availability_topic.clone(),
                    "device": context.full.mqtt.discovery_device_object.clone(),
                    "force_update": true,
                    "icon": options.icon,
                    "name": options.display_name,
                    "suggested_display_precision": options.suggested_precision,
                    "state_class": "measurement",
                    "state_topic": state_topic.as_str(),
                    "object_id": entity_id.as_str(),
                    "unique_id": entity_id.as_str()
                };

                if let Some(unit) = options.unit {
                    object.insert("unit_of_measurement", unit).unwrap();
                }

                object
            }),
        )
        .await?;

    let mqtt_client = context.full.mqtt.client.clone();

    tokio::spawn(async move {
        let mut interval = tokio::time::interval(interval);
        interval.set_missed_tick_behavior(MissedTickBehavior::Delay);

        loop {
            interval.tick().await;

            let result = fetch();

            match result {
                Err(error) => log::error!("{:#}", anyhow!(error)),

                Ok(value) => {
                    if let Err(error) = mqtt_client.publish(&state_topic, QoS::ExactlyOnce, true, value).await {
                        log::error!("{:#}", anyhow!(error));
                    };
                }
            }
        }
    });

    Ok(())
}