launcher/src/power.rs

//! `power on` / `power off` subcommands.
//!
//! Power-off: SSH `sudo shutdown now` to each node → wait
//! `[home_assistant].shutdown_delay_seconds` (default 30s) → POST
//! `turn_off` to Home Assistant for each per-node `entity` and each
//! cluster-level `extra_entities`.
//!
//! Power-on: just POST `turn_on` to Home Assistant for each per-node
//! `entity` and each cluster-level `extra_entities` (the smart plug
//! handles waking the machine).

use anyhow::Result;
use std::time::Duration;
use tracing::{info, warn};

use crate::config::{Cluster, Config, Node};
use crate::home_assistant::Client as HaClient;

pub mod plug;
pub mod toggle;
pub use plug::{OnAction, decide_on, stale_nodes};
use toggle::{call, extra_ids, off_extras, toggle_entities, toggle_one};
mod mikro;

/// Flow for `power off`: SSH shutdown + wait + HA `turn_off`.
pub async fn power_off(
    config: &Config,
    cluster: &Cluster,
    node_names: &[String],
    dry_run: bool,
) -> Result<()> {
    let targets = select_targets(cluster, node_names, "power off")?;

    // Cluster-level extras (the fabric switch, a cooling fan) are SHARED by
    // every node in the cluster.  Powering off a SUBSET (`--node`) must not
    // take the fabric down for the nodes that stay up — that wedges live
    // multi-node serves (NCCL over the switch dies).  And even a whole-cluster
    // off must keep any extra that ANOTHER cluster with an open run references
    // (the fabric switch is shared across clusters — powering `8nodes-switch`
    // down must not kill a serve on `dgx58-switch` over the same switch).
    let held = toggle::extras_held_by_other_clusters(config, cluster);
    let held_ids: std::collections::HashSet<String> = held.keys().cloned().collect();
    let extras = off_extras(cluster, node_names, &held_ids);
    if !node_names.is_empty() && !cluster.extra_entities.is_empty() {
        info!(
            "{}",
            fluent_i18n::t!("power-off-extras-skipped", {
                "nodes" => node_names.join(", "),
                "extras" => extra_ids(cluster).join(", ")
            })
        );
    }
    if node_names.is_empty() {
        for e in extra_ids(cluster) {
            if let Some(clusters) = held.get(&e) {
                info!(
                    "{}",
                    fluent_i18n::t!("power-off-extras-held", {
                        "entity" => e.clone(),
                        "clusters" => clusters.join(", ")
                    })
                );
            }
        }
    }

    if dry_run {
        warn!("{}", fluent_i18n::t!("power-off-dryrun"));
        mikro::dry_off(config, &targets, &extras);
        return Ok(());
    }

    // 1) SSH shutdown every target in parallel.
    // NOTE: This is fire-and-forget — if a node is unreachable, its
    // SSH failure is logged but the flow continues.  The HA `turn_off`
    // still fires for all targets regardless of SSH success.  This is
    // by design: already-offline nodes get their HA entity toggled off,
    // which is idempotent.
    info!(
        "{}",
        fluent_i18n::t!("power-off-head", { "nodes" => targets.len() })
    );
    let mut handles = Vec::with_capacity(targets.len());
    for node in targets.iter().cloned() {
        let config = config.clone();
        handles.push(crate::locale::spawn_with_locale(async move {
            let cmd = "sudo shutdown now".to_string();
            let r = crate::ssh::run_ssh(&config, &node, cmd.clone(), false).await;
            (node, cmd, r)
        }));
    }

    let mut any_ssh_failed = false;
    for h in handles {
        let (node, cmd, r) = h.await?;
        if let Err(e) = r {
            any_ssh_failed = true;
            warn!(
                "{}",
                fluent_i18n::t!("power-off-shutdown-failed", { "name" => &node.name, "cmd" => &cmd, "error" => e.to_string() })
            );
        }
    }

    // 2) Wait so the OS actually shuts down before we cut power.
    let delay = config
        .home_assistant
        .as_ref()
        .map(|c| c.shutdown_delay_seconds)
        .unwrap_or(30);
    info!(
        "{}",
        fluent_i18n::t!("power-off-waiting", { "secs" => delay })
    );
    tokio::time::sleep(Duration::from_secs(delay)).await;

    // 3) HA turn_off: per-node entities, then cluster-level extras.
    //    Mikrotik names are skipped here — they get a graceful RouterOS
    //    API shutdown in step 4 (after HA, so the router dies last).
    let mikrotik_off: Vec<String> = if let Some(ha_cfg) = &config.home_assistant {
        let ha = HaClient::new(ha_cfg)?;
        crate::home_assistant::warn_if_token_looks_short(&ha_cfg.token);
        toggle_entities(
            &ha,
            "turn_off",
            config.mikrotik.as_ref(),
            &targets,
            &extras,
            "cluster",
        )
        .await
    } else {
        info!("{}", fluent_i18n::t!("power-off-no-ha"));
        // No HA at all — still shut down any mikrotik device referenced.
        mikro::collect_off(config, &targets, &extras)
    };

    // 4) RouterOS graceful shutdown for each referenced device.
    mikro::shutdown_named(config, &mikrotik_off).await;

    if any_ssh_failed {
        warn!("{}", fluent_i18n::t!("power-off-ssh-failures"));
    }
    Ok(())
}

/// Flow for `power on`: HA `turn_on` for each node with an `entity`
/// and each cluster-level `extra_entities`.  No SSH needed — the smart
/// plug handles waking the machine.
/// Flow for `power on`: HA `turn_on` for each node with an `entity`
/// Flow for `power on`: per-node plug decisions (`decide_on` + the caller's
/// `stale` list), then cluster-level `extra_entities`.  Returns one report
/// line per node — the job surface (web job message, CLI log) tells what
/// actually happened instead of a blanket "ok".  No SSH needed — the smart
/// plug (plus a cycle for stale-on plugs) wakes the machine.
pub async fn power_on(
    config: &Config,
    cluster: &Cluster,
    node_names: &[String],
    dry_run: bool,
    stale: &[String],
) -> Result<Vec<String>> {
    let targets = select_targets(cluster, node_names, "power on")?;

    let Some(ha_cfg) = &config.home_assistant else {
        anyhow::bail!(
            "[home_assistant] section is required for `power on` \
             (the launcher has no other way to wake a node)"
        );
    };
    let ha = HaClient::new(ha_cfg)?;
    crate::home_assistant::warn_if_token_looks_short(&ha_cfg.token);

    if dry_run {
        warn!("{}", fluent_i18n::t!("power-on-dryrun"));
        mikro::dry_on(config, &targets, &extra_ids(cluster));
        return Ok(Vec::new());
    }

    info!(
        "{}",
        fluent_i18n::t!("power-on-head", { "nodes" => targets.len() })
    );
    let mut report: Vec<String> = Vec::new();
    for node in &targets {
        let Some(entity) = &node.entity else {
            warn!(
                "{}",
                fluent_i18n::t!("ha-skip-no-entity", { "name" => &node.name })
            );
            report.push(fluent_i18n::t!("power-on-skip", { "name" => &node.name }));
            continue;
        };
        if mikro::device(config.mikrotik.as_ref(), entity).is_some() {
            // RouterOS devices are not HA plugs; wake via the relay as before
            toggle_one(
                &ha,
                "turn_on",
                config.mikrotik.as_ref(),
                entity,
                &node.name,
                &mut Vec::new(),
            )
            .await;
            report.push(fluent_i18n::t!("power-on-on", { "name" => &node.name }));
            continue;
        }
        let action = decide_on(
            ha.get_state(entity).await.as_deref(),
            stale.contains(&node.name),
        );
        match action {
            OnAction::TurnOn => {
                if let Err(e) = call(&ha, entity, "turn_on").await {
                    warn!(
                        "{}",
                        fluent_i18n::t!("ha-call-failed", { "service" => "turn_on", "entity" => entity.clone(), "error" => e.to_string() })
                    );
                    report.push(fluent_i18n::t!("power-on-failed", { "name" => &node.name, "error" => e.to_string() }));
                } else {
                    report.push(fluent_i18n::t!("power-on-on", { "name" => &node.name }));
                }
            }
            OnAction::Cycle => {
                warn!(
                    "{}",
                    fluent_i18n::t!("power-on-cycling", { "name" => &node.name, "entity" => entity.clone() })
                );
                let r1 = call(&ha, entity, "turn_off").await;
                // the relay opens asynchronously — wait until HA actually
                // reports off (≤10 s), then a settle dwell, then re-apply
                let mut saw_off = false;
                for _ in 0..20 {
                    tokio::time::sleep(Duration::from_millis(500)).await;
                    if ha.get_state(entity).await.as_deref() != Some("on") {
                        saw_off = true;
                        break;
                    }
                }
                if !saw_off {
                    warn!(
                        "{}",
                        fluent_i18n::t!("power-on-off-wait-timeout", { "name" => &node.name, "entity" => entity.clone() })
                    );
                }
                // the PDU needs a real outage: a 2-3 s cut left the PSU
                // rail energized and the machine stayed in soft-off — 30 s
                // reliably re-triggers the power-on
                tokio::time::sleep(Duration::from_secs(30)).await;
                let r2 = call(&ha, entity, "turn_on").await;
                match (r1, r2) {
                    (Ok(()), Ok(())) => {
                        report.push(fluent_i18n::t!("power-on-cycled", { "name" => &node.name }))
                    }
                    (e1, e2) => {
                        let e = e1
                            .err()
                            .or(e2.err())
                            .unwrap_or_else(|| anyhow::anyhow!("cycle"));
                        report.push(fluent_i18n::t!("power-on-failed", { "name" => &node.name, "error" => e.to_string() }));
                    }
                }
            }
            OnAction::AlreadyOn => {
                warn!(
                    "{}",
                    fluent_i18n::t!("power-on-was-on", { "name" => &node.name })
                );
                report.push(fluent_i18n::t!("power-on-was-on", { "name" => &node.name }));
            }
        }
    }
    // cluster-level extras (fan / relay / fabric switch) as before
    let _ = toggle_entities(
        &ha,
        "turn_on",
        config.mikrotik.as_ref(),
        &[],
        &extra_ids(cluster),
        "cluster",
    )
    .await;
    Ok(report)
}

/// Flow for `power switch --from <src> --to <tgt>`:
/// 1. Shut down + HA turn_off nodes/entities unique to `from`,
/// 2. HA turn_on all nodes/entities in `to`.
pub async fn power_switch(
    config: &Config,
    from: &Cluster,
    to: &Cluster,
    dry_run: bool,
) -> Result<()> {
    let to_names: std::collections::HashSet<&str> =
        to.nodes.iter().map(|n| n.name.as_str()).collect();
    let disable_nodes: Vec<Node> = from
        .nodes
        .iter()
        .filter(|n| !to_names.contains(n.name.as_str()))
        .cloned()
        .collect();
    let to_extras: std::collections::HashSet<&str> =
        to.extra_entities.iter().map(|e| e.entity()).collect();
    let disable_extras: Vec<String> = from
        .extra_entities
        .iter()
        .filter(|e| !to_extras.contains(e.entity()))
        .map(|e| e.entity().to_string())
        .collect();

    info!(
        "{}",
        fluent_i18n::t!("power-switch-head", { "from" => &from.name, "to" => &to.name, "nodes" => disable_nodes.len(), "entities" => disable_extras.len() })
    );

    if dry_run {
        mikro::dry_switch_off(config, &disable_nodes, &disable_extras);
        mikro::dry_switch_on(config, &to.nodes, &extra_ids(to));
        return Ok(());
    }

    // SSH shutdown nodes not in target
    let mut handles = Vec::with_capacity(disable_nodes.len());
    for node in disable_nodes.iter().cloned() {
        let c = config.clone();
        handles.push(crate::locale::spawn_with_locale(async move {
            (
                node.name.clone(),
                crate::ssh::run_ssh(&c, &node, "sudo shutdown now".into(), false).await,
            )
        }));
    }
    for h in handles {
        let (name, r) = h.await?;
        if let Err(e) = r {
            warn!(
                "{}",
                fluent_i18n::t!("power-switch-shutdown-failed", { "name" => &name, "error" => e.to_string() })
            );
        }
    }

    let delay = config
        .home_assistant
        .as_ref()
        .map(|c| c.shutdown_delay_seconds)
        .unwrap_or(30);
    if !disable_nodes.is_empty() {
        info!(
            "{}",
            fluent_i18n::t!("power-switch-waiting", { "secs" => delay })
        );
        tokio::time::sleep(Duration::from_secs(delay)).await;
    } else {
        // No nodes to shut down — only extra_entities differ.  HA
        // turn_off / turn_on runs immediately without delay.
    }

    if let Some(ha_cfg) = &config.home_assistant {
        let ha = HaClient::new(ha_cfg)?;
        crate::home_assistant::warn_if_token_looks_short(&ha_cfg.token);

        // Disable phase: per-node entities + extras unique to `from`.
        // Routed through `toggle_entities` (which validates + logs failures)
        // instead of ad-hoc loops with `let _ =` (which swallowed errors)
        // and dead `if dry_run` branches (dry-run returned at the top).
        let mikrotik_off = toggle_entities(
            &ha,
            "turn_off",
            config.mikrotik.as_ref(),
            &disable_nodes,
            &disable_extras,
            "from",
        )
        .await;

        // RouterOS devices named above: graceful API shutdown (after HA).
        mikro::shutdown_named(config, &mikrotik_off).await;

        info!(
            "{}",
            fluent_i18n::t!("power-switch-powering-on", { "nodes" => to.nodes.len(), "entities" => to.extra_entities.len(), "to" => &to.name })
        );
        let _ = toggle_entities(
            &ha,
            "turn_on",
            config.mikrotik.as_ref(),
            &to.nodes,
            &extra_ids(to),
            "to",
        )
        .await;
    } else {
        info!("{}", fluent_i18n::t!("power-switch-no-ha"));
        // No HA, but mikrotik devices named in `from` still get the
        // graceful API shutdown.
        mikro::shutdown_named(
            config,
            &mikro::collect_off(config, &disable_nodes, &disable_extras),
        )
        .await;
    }
    Ok(())
}

pub(crate) fn select_targets(
    cluster: &Cluster,
    node_names: &[String],
    op: &str,
) -> Result<Vec<Node>> {
    let targets: Vec<Node> = if node_names.is_empty() {
        cluster.nodes.clone()
    } else {
        let mut found: Vec<Node> = Vec::new();
        let mut missing: Vec<String> = Vec::new();
        for nn in node_names {
            if let Some(node) = cluster.nodes.iter().find(|n| &n.name == nn) {
                found.push(node.clone());
            } else {
                missing.push(nn.clone());
            }
        }
        if !missing.is_empty() {
            anyhow::bail!(
                "{}: node(s) {} not found in cluster '{}'",
                op,
                missing.join(", "),
                cluster.name
            );
        }
        found
    };
    if targets.is_empty() {
        anyhow::bail!("{}: cluster '{}' has no nodes", op, cluster.name);
    }
    Ok(targets)
}

#[cfg(test)]
mod tests;