consortium/

propagation.rs

//! Tree propagation engine.
//!
//! Rust implementation of `ClusterShell.Propagation`.
//!
//! This module provides:
//! - [`PropagationTreeRouter`]: Routes commands through the topology tree by
//!   computing next-hop gateways and managing fanout/load-balancing.
//! - [`PropagationChannel`]: State machine for gateway-to-gateway channel
//!   lifecycle (configuration, control, gather phases).

use std::collections::{HashMap, VecDeque};

use crate::communication::{Channel, Message, MessageProcessingError};
use crate::node_set::NodeSet;
use crate::topology::{TopologyError, TopologyTree};

use thiserror::Error;

// ============================================================================
// Errors
// ============================================================================

/// Errors raised during route resolving operations.
#[derive(Error, Debug)]
pub enum RouteResolvingError {
    /// Invalid root or gateway node.
    #[error("Invalid root or gateway node: {0}")]
    InvalidRoot(String),

    /// No route available to the destination.
    #[error("No route available to {0}")]
    NoRoute(String),

    /// Destination host is unreachable.
    #[error("Invalid destination: {0}, host is unreachable")]
    Unreachable(String),

    /// Source equals destination.
    #[error("Invalid resolution request: {0} -> {1}")]
    SameSourceDest(String, String),

    /// Topology error.
    #[error("Topology error: {0}")]
    Topology(#[from] TopologyError),

    /// NodeSet error.
    #[error("NodeSet error: {0}")]
    NodeSet(#[from] crate::node_set::NodeSetError),
}

// ============================================================================
// PropagationTreeRouter
// ============================================================================

/// A routing table entry: (reachable network, next-hop gateways).
#[derive(Debug, Clone)]
pub struct RouteEntry {
    /// Set of all nodes reachable through this route (transitive children).
    network: NodeSet,
    /// Direct next-hop gateway(s) for this route.
    nexthops: NodeSet,
}

/// Performs route resolving operations within a propagation tree.
///
/// The router parses the topology tree to generate a routing table, then
/// provides [`next_hop()`](PropagationTreeRouter::next_hop) and
/// [`dispatch()`](PropagationTreeRouter::dispatch) methods for directing
/// messages to the correct gateway.
///
/// Load balancing is achieved by tracking the number of connections
/// (fanin) per gateway node.
#[derive(Debug)]
pub struct PropagationTreeRouter {
    /// The root node of this router.
    root: String,
    /// Fanout setting (0 = no limit).
    fanout: usize,
    /// Connections count per gateway node (for load balancing).
    nodes_fanin: HashMap<String, usize>,
    /// Routing table: list of (network, nexthops).
    table: Vec<RouteEntry>,
    /// Set of nodes marked as unreachable.
    unreachable_hosts: NodeSet,
}

impl PropagationTreeRouter {
    /// Create a new router from a root node and topology tree.
    ///
    /// The routing table is generated by traversing the tree from the root
    /// node's position downward.
    ///
    /// # Errors
    /// Returns [`RouteResolvingError::InvalidRoot`] if the root node cannot
    /// be found in the topology.
    pub fn new(
        root: &str,
        topology: &TopologyTree,
        fanout: usize,
    ) -> Result<Self, RouteResolvingError> {
        let mut router = Self {
            root: root.to_string(),
            fanout,
            nodes_fanin: HashMap::new(),
            table: Vec::new(),
            unreachable_hosts: NodeSet::new(),
        };
        router.table_generate(root, topology)?;
        Ok(router)
    }

    /// Generate the routing table from the topology tree.
    ///
    /// For each direct child group of the root, compute the set of all
    /// transitively reachable nodes (the "network") and store it alongside
    /// the next-hop gateway nodeset.
    fn table_generate(
        &mut self,
        root: &str,
        topology: &TopologyTree,
    ) -> Result<(), RouteResolvingError> {
        // Validate the root exists in the topology
        let _root_group = topology
            .find_nodegroup(root)
            .map_err(|_| RouteResolvingError::InvalidRoot(root.to_string()))?;

        self.table.clear();

        // We need to find the root TopologyNodeGroup and enumerate its
        // direct children. Each child group becomes a route entry, with
        // its "network" being all transitively reachable nodes below it.
        let root_tng = topology
            .root()
            .ok_or_else(|| RouteResolvingError::InvalidRoot(root.to_string()))?;

        // The root of the topology tree should contain our root node.
        // If the root_tng doesn't contain our node, we need to find
        // the subtree rooted at our node.
        let subtree_root = Self::find_group_containing(root_tng, root)
            .ok_or_else(|| RouteResolvingError::InvalidRoot(root.to_string()))?;

        for child in subtree_root.children() {
            let mut network = NodeSet::new();
            Self::collect_reachable(child, &mut network);
            self.table.push(RouteEntry {
                network,
                nexthops: child.nodeset.clone(),
            });
        }

        Ok(())
    }

    /// Find the TopologyNodeGroup that contains the given node name.
    fn find_group_containing<'a>(
        group: &'a crate::topology::TopologyNodeGroup,
        node: &str,
    ) -> Option<&'a crate::topology::TopologyNodeGroup> {
        if group.nodeset.contains(node) {
            return Some(group);
        }
        for child in group.children() {
            if let Some(found) = Self::find_group_containing(child, node) {
                return Some(found);
            }
        }
        None
    }

    /// Recursively collect all nodes reachable below a group (its children,
    /// grandchildren, etc.) into `dest`. Does NOT include the group itself.
    fn collect_reachable(group: &crate::topology::TopologyNodeGroup, dest: &mut NodeSet) {
        for child in group.children() {
            dest.update(&child.nodeset);
            Self::collect_reachable(child, dest);
        }
    }

    /// Dispatch nodes from a target nodeset to directly connected gateways.
    ///
    /// Acts as an iterator, yielding `(gateway_nodeset, target_nodes)` pairs.
    /// Provides load balancing between gateways. Remaining nodes that don't
    /// match any route are treated as directly connected.
    pub fn dispatch(&self, dst: &NodeSet) -> Vec<(NodeSet, NodeSet)> {
        let mut remaining = dst.clone();
        let mut results = Vec::new();

        // Check for remote targets requiring a gateway
        for entry in &self.table {
            let dst_inter = entry.network.intersection(&remaining);
            if !dst_inter.is_empty() {
                remaining.difference_update(&dst_inter);
                // For each host in the intersection, find best next hop
                for host_name in dst_inter.iter() {
                    if let Ok(hop) = self.next_hop_str(&host_name) {
                        let mut hop_ns = NodeSet::new();
                        let _ = hop_ns.update_str(&hop);
                        let mut host_ns = NodeSet::new();
                        let _ = host_ns.update_str(&host_name);
                        results.push((hop_ns, host_ns));
                    }
                }
            }
        }

        // Remaining nodes are directly connected
        if !remaining.is_empty() {
            results.push((remaining.clone(), remaining));
        }

        results
    }

    /// Perform next-hop resolution for a single destination node.
    ///
    /// Returns the best gateway to use based on current load (connections
    /// tracked in `nodes_fanin`).
    pub fn next_hop(&mut self, dst: &NodeSet) -> Result<NodeSet, RouteResolvingError> {
        let dst_str = dst.to_string();

        // Check unreachable
        if !self.unreachable_hosts.intersection(dst).is_empty() {
            return Err(RouteResolvingError::Unreachable(dst_str));
        }

        // Can't resolve if source == destination
        if self.root == dst_str {
            return Err(RouteResolvingError::SameSourceDest(
                self.root.clone(),
                dst_str,
            ));
        }

        for entry in &self.table {
            // destination contained in current network
            if !entry.network.intersection(dst).is_empty() {
                let res = self.best_next_hop(&entry.nexthops)?;
                let res_str = res.to_string();
                *self.nodes_fanin.entry(res_str).or_insert(0) += dst.len();
                return Ok(res);
            }
            // destination is a direct next-hop
            if !entry.nexthops.intersection(dst).is_empty() {
                return Ok(dst.clone());
            }
        }

        Err(RouteResolvingError::NoRoute(dst_str))
    }

    /// String-based next_hop for convenience (non-mutating for dispatch).
    fn next_hop_str(&self, dst: &str) -> Result<String, RouteResolvingError> {
        // Check unreachable
        if self.unreachable_hosts.contains(dst) {
            return Err(RouteResolvingError::Unreachable(dst.to_string()));
        }

        if self.root == dst {
            return Err(RouteResolvingError::SameSourceDest(
                self.root.clone(),
                dst.to_string(),
            ));
        }

        let dst_ns = NodeSet::parse(dst)?;

        for entry in &self.table {
            if !entry.network.intersection(&dst_ns).is_empty() {
                let res = self.best_next_hop(&entry.nexthops)?;
                return Ok(res.to_string());
            }
            if !entry.nexthops.intersection(&dst_ns).is_empty() {
                return Ok(dst.to_string());
            }
        }

        Err(RouteResolvingError::NoRoute(dst.to_string()))
    }

    /// Mark a node as unreachable.
    pub fn mark_unreachable(&mut self, dst: &str) {
        let _ = self.unreachable_hosts.update_str(dst);
    }

    /// Find the best next-hop from a set of candidates.
    ///
    /// Selects the gateway with the fewest current connections (load balancing).
    fn best_next_hop(&self, candidates: &NodeSet) -> Result<NodeSet, RouteResolvingError> {
        let available = candidates.difference(&self.unreachable_hosts);
        if available.is_empty() {
            return Err(RouteResolvingError::NoRoute(
                "all candidates unreachable".into(),
            ));
        }

        let mut best: Option<String> = None;
        let mut best_connections = usize::MAX;

        for host in available.iter() {
            let connections = self.nodes_fanin.get(&host).copied().unwrap_or(0);
            if connections < best_connections {
                best_connections = connections;
                best = Some(host);
            }
        }

        match best {
            Some(host) => {
                let mut ns = NodeSet::new();
                let _ = ns.update_str(&host);
                Ok(ns)
            }
            None => Err(RouteResolvingError::NoRoute(
                "no candidate available".into(),
            )),
        }
    }

    /// Get the root node name.
    pub fn root(&self) -> &str {
        &self.root
    }

    /// Get the fanout setting.
    pub fn fanout(&self) -> usize {
        self.fanout
    }

    /// Get the routing table for inspection.
    pub fn table(&self) -> &[RouteEntry] {
        &self.table
    }
}

// ============================================================================
// PropagationChannel
// ============================================================================

/// State of a PropagationChannel's finite state machine.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChannelState {
    /// Initial state — channel not yet opened.
    Init,
    /// Configuration sent, waiting for ACK.
    Config,
    /// Control message sent, waiting for ACK.
    Control,
    /// Gathering results from the subtree.
    Gather,
    /// Channel is closed.
    Closed,
}

/// Admin-node propagation channel.
///
/// Manages the lifecycle of communication with a gateway node:
/// 1. **Init**: Send configuration (serialized topology).
/// 2. **Config**: Wait for ACK, then send control message (command).
/// 3. **Control**: Wait for ACK, then enter gather state.
/// 4. **Gather**: Collect results (stdout, stderr, retcodes, timeouts, routing).
///
/// This is a state machine that drives the [`Channel`] transport.
#[derive(Debug)]
pub struct PropagationChannel {
    /// The underlying communication channel.
    pub channel: Channel,
    /// The gateway node this channel connects to.
    pub gateway: String,
    /// Current FSM state.
    pub state: ChannelState,
    /// Queue of control messages waiting to be sent.
    send_queue: VecDeque<Message>,
    /// Last return code from the gateway process.
    pub last_rc: Option<i32>,
    /// Write history for tracking ACKs to write operations.
    /// Each entry: (msgid, nodes, byte_count).
    write_history: VecDeque<(u64, String, usize)>,
    /// Collected results: (node, stream_name, data).
    pub results: Vec<PropagationResult>,
}

/// A result received from the propagation subtree.
#[derive(Debug, Clone)]
pub enum PropagationResult {
    /// Standard output line from a node.
    StdOut { node: String, data: Vec<u8> },
    /// Standard error line from a node.
    StdErr { node: String, data: Vec<u8> },
    /// Return code from a node.
    Retcode { node: String, rc: i32 },
    /// Timeout for a node.
    Timeout { node: String },
    /// Routing event.
    Routing {
        event: String,
        gateway: String,
        targets: String,
    },
}

impl PropagationChannel {
    /// Create a new propagation channel for a gateway.
    pub fn new(gateway: &str) -> Self {
        Self {
            channel: Channel::new(true), // initiator
            gateway: gateway.to_string(),
            state: ChannelState::Init,
            send_queue: VecDeque::new(),
            last_rc: None,
            write_history: VecDeque::new(),
            results: Vec::new(),
        }
    }

    /// Start the propagation channel — send configuration.
    ///
    /// Serializes the topology and sends it as a ConfigurationMessage.
    pub fn start(&mut self, topology_bytes: &[u8]) {
        // Open the channel
        self.channel.opened = true;

        // Send configuration message
        let mut cfg = Message::configuration(&self.gateway);
        cfg.data_encode(topology_bytes);
        self.channel.send(&cfg);
        self.state = ChannelState::Config;
    }

    /// Enqueue a control message for sending.
    ///
    /// If the channel is already set up and the queue is empty, the message
    /// is sent immediately. Otherwise it's queued for later.
    pub fn send_queued(&mut self, msg: Message) {
        if self.channel.setup && self.send_queue.is_empty() {
            self.channel.send(&msg);
        } else {
            self.send_queue.push_back(msg);
        }
    }

    /// Send one queued message, if any.
    fn send_dequeue(&mut self) {
        if let Some(msg) = self.send_queue.pop_front() {
            self.channel.send(&msg);
        }
    }

    /// Create and enqueue a shell command for execution through the channel.
    pub fn shell(&mut self, nodes: &str, command: &[u8], worker_id: u64) {
        let mut ctl = Message::control(worker_id);
        if let Message::Control {
            ref mut action,
            ref mut target,
            ..
        } = ctl
        {
            *action = "shell".to_string();
            *target = nodes.to_string();
        }
        ctl.data_encode(command);
        self.send_queued(ctl);
    }

    /// Create and enqueue a write operation through the channel.
    pub fn write(&mut self, nodes: &str, buf: &[u8], worker_id: u64) {
        let mut ctl = Message::control(worker_id);
        if let Message::Control {
            ref mut action,
            ref mut target,
            ..
        } = ctl
        {
            *action = "write".to_string();
            *target = nodes.to_string();
        }
        let msgid = ctl.msgid();
        ctl.data_encode(buf);
        self.write_history
            .push_back((msgid, nodes.to_string(), buf.len()));
        self.send_queued(ctl);
    }

    /// Create and enqueue an EOF operation through the channel.
    pub fn set_write_eof(&mut self, nodes: &str, worker_id: u64) {
        let mut ctl = Message::control(worker_id);
        if let Message::Control {
            ref mut action,
            ref mut target,
            ..
        } = ctl
        {
            *action = "eof".to_string();
            *target = nodes.to_string();
        }
        self.send_queued(ctl);
    }

    /// Process an incoming message according to the current FSM state.
    ///
    /// Returns any propagation results generated by this message.
    pub fn recv(&mut self, msg: Message) -> Result<Vec<PropagationResult>, MessageProcessingError> {
        let mut new_results = Vec::new();

        match msg.ident() {
            "END" => {
                self.state = ChannelState::Closed;
                return Ok(new_results);
            }
            "CHA" => {
                self.channel.opened = true;
                return Ok(new_results);
            }
            _ => {}
        }

        match self.state {
            ChannelState::Config => {
                // Waiting for ACK to configuration
                if msg.ident() == "ACK" {
                    self.channel.setup = true;
                    self.state = ChannelState::Gather;
                    self.send_dequeue();
                }
            }
            ChannelState::Control | ChannelState::Gather => {
                match msg.ident() {
                    "ACK" => {
                        // Check if ACK matches a write history entry
                        if let Message::Ack { ack, .. } = &msg {
                            if let Some(back) = self.write_history.back() {
                                if back.0 == *ack {
                                    let _ = self.write_history.pop_back();
                                }
                            }
                        }
                        self.send_dequeue();
                    }
                    "OUT" => {
                        if let Some(nodes_str) = msg.nodes() {
                            if let Ok(decoded) = msg.data_decode() {
                                let ns = NodeSet::parse(nodes_str).unwrap_or_default();
                                for node in ns.iter() {
                                    let result = PropagationResult::StdOut {
                                        node,
                                        data: decoded.clone(),
                                    };
                                    new_results.push(result);
                                }
                            }
                        }
                    }
                    "SER" => {
                        if let Some(nodes_str) = msg.nodes() {
                            if let Ok(decoded) = msg.data_decode() {
                                let ns = NodeSet::parse(nodes_str).unwrap_or_default();
                                for node in ns.iter() {
                                    let result = PropagationResult::StdErr {
                                        node,
                                        data: decoded.clone(),
                                    };
                                    new_results.push(result);
                                }
                            }
                        }
                    }
                    "RET" => {
                        if let Message::Retcode { retcode, nodes, .. } = &msg {
                            let ns = NodeSet::parse(nodes).unwrap_or_default();
                            for node in ns.iter() {
                                new_results.push(PropagationResult::Retcode { node, rc: *retcode });
                            }
                        }
                    }
                    "TIM" => {
                        if let Some(nodes_str) = msg.nodes() {
                            let ns = NodeSet::parse(nodes_str).unwrap_or_default();
                            for node in ns.iter() {
                                new_results.push(PropagationResult::Timeout { node });
                            }
                        }
                    }
                    "RTR" => {
                        if let Message::Routing {
                            event,
                            gateway,
                            targets,
                            ..
                        } = &msg
                        {
                            new_results.push(PropagationResult::Routing {
                                event: event.clone(),
                                gateway: gateway.clone(),
                                targets: targets.clone(),
                            });
                        }
                    }
                    "ERR" => {
                        if let Message::Error { reason, .. } = &msg {
                            return Err(MessageProcessingError::ParseError(format!(
                                "{}: {}",
                                self.gateway, reason
                            )));
                        }
                    }
                    _ => {
                        // Ignore unknown messages in gather state
                    }
                }
            }
            _ => {}
        }

        self.results.extend(new_results.clone());
        Ok(new_results)
    }

    /// Handle the gateway process closing.
    pub fn ev_hup(&mut self, rc: i32) {
        self.last_rc = Some(rc);
    }

    /// Handle the channel closing.
    pub fn ev_close(&mut self) {
        self.state = ChannelState::Closed;
    }

    /// Get the current channel state.
    pub fn current_state(&self) -> ChannelState {
        self.state
    }

    /// Check if the channel is in gather state (ready for results).
    pub fn is_gathering(&self) -> bool {
        self.state == ChannelState::Gather
    }

    /// Check if the channel is closed.
    pub fn is_closed(&self) -> bool {
        self.state == ChannelState::Closed
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::topology::{TopologyGraph, TopologyTree};

    /// Helper to build a simple topology tree for testing.
    /// Structure: admin -> gw1 -> node[1-4], gw2 -> node[5-8]
    fn build_test_topology() -> TopologyTree {
        let mut graph = TopologyGraph::new();
        let admin = NodeSet::parse("admin").unwrap();
        let gw1 = NodeSet::parse("gw1").unwrap();
        let gw2 = NodeSet::parse("gw2").unwrap();
        let nodes1 = NodeSet::parse("node[1-4]").unwrap();
        let nodes2 = NodeSet::parse("node[5-8]").unwrap();

        graph.add_route(admin.clone(), gw1.clone()).unwrap();
        graph.add_route(admin, gw2.clone()).unwrap();
        graph.add_route(gw1, nodes1).unwrap();
        graph.add_route(gw2, nodes2).unwrap();
        graph.to_tree("admin").unwrap()
    }

    // -- RouteResolvingError tests ------------------------------------------

    #[test]
    fn test_error_display() {
        let err = RouteResolvingError::InvalidRoot("badnode".into());
        assert!(err.to_string().contains("badnode"));

        let err = RouteResolvingError::NoRoute("node99".into());
        assert!(err.to_string().contains("node99"));

        let err = RouteResolvingError::Unreachable("gw1".into());
        assert!(err.to_string().contains("gw1"));

        let err = RouteResolvingError::SameSourceDest("a".into(), "a".into());
        assert!(err.to_string().contains("a"));
    }

    // -- PropagationTreeRouter tests ----------------------------------------

    #[test]
    fn test_router_creation() {
        let tree = build_test_topology();
        let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
        assert_eq!(router.root(), "admin");
        assert_eq!(router.fanout(), 0);
    }

    #[test]
    fn test_router_invalid_root() {
        let tree = build_test_topology();
        let result = PropagationTreeRouter::new("nonexistent", &tree, 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_router_next_hop() {
        let tree = build_test_topology();
        let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();

        // node1 should be reachable through gw1
        let dst = NodeSet::parse("node1").unwrap();
        let hop = router.next_hop(&dst).unwrap();
        assert_eq!(hop.to_string(), "gw1");

        // node5 should be reachable through gw2
        let dst = NodeSet::parse("node5").unwrap();
        let hop = router.next_hop(&dst).unwrap();
        assert_eq!(hop.to_string(), "gw2");
    }

    #[test]
    fn test_router_same_source_dest() {
        let tree = build_test_topology();
        let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
        let dst = NodeSet::parse("admin").unwrap();
        let result = router.next_hop(&dst);
        assert!(result.is_err());
    }

    #[test]
    fn test_router_mark_unreachable() {
        let tree = build_test_topology();
        let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();

        router.mark_unreachable("gw1");

        // node1 goes through gw1, but gw1 is unreachable
        // The best_next_hop should fail since the only candidate is unreachable
        let dst = NodeSet::parse("node1").unwrap();
        let result = router.next_hop(&dst);
        assert!(result.is_err());
    }

    #[test]
    fn test_router_dispatch_directly_connected() {
        let tree = build_test_topology();
        let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();

        // Dispatch nodes not in any route -> treated as directly connected
        let dst = NodeSet::parse("localnode").unwrap();
        let results = router.dispatch(&dst);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0.to_string(), "localnode");
    }

    #[test]
    fn test_router_dispatch_remote() {
        let tree = build_test_topology();
        let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();

        let dst = NodeSet::parse("node[1-2]").unwrap();
        let results = router.dispatch(&dst);
        // Should dispatch through gw1
        assert!(!results.is_empty());
        for (hop, _target) in &results {
            assert_eq!(hop.to_string(), "gw1");
        }
    }

    #[test]
    fn test_router_load_balancing() {
        // Build a topology with two gateways to same network
        // admin -> gw[1-2] -> node[1-8]
        let mut graph = TopologyGraph::new();
        let admin = NodeSet::parse("admin").unwrap();
        let gws = NodeSet::parse("gw[1-2]").unwrap();
        let nodes = NodeSet::parse("node[1-8]").unwrap();
        graph.add_route(admin, gws).unwrap();

        // Each gw gets half the nodes
        let gw1 = NodeSet::parse("gw1").unwrap();
        let gw2 = NodeSet::parse("gw2").unwrap();
        graph
            .add_route(gw1, NodeSet::parse("node[1-4]").unwrap())
            .unwrap();
        graph
            .add_route(gw2, NodeSet::parse("node[5-8]").unwrap())
            .unwrap();

        let tree = graph.to_tree("admin").unwrap();
        let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();

        // Route several nodes — should distribute across gateways
        let d1 = NodeSet::parse("node1").unwrap();
        let _h1 = router.next_hop(&d1).unwrap();
        let d2 = NodeSet::parse("node5").unwrap();
        let _h2 = router.next_hop(&d2).unwrap();

        // Both should resolve to different gateways
        // (node1 -> gw1, node5 -> gw2 based on topology)
    }

    // -- PropagationChannel tests -------------------------------------------

    #[test]
    fn test_channel_creation() {
        let ch = PropagationChannel::new("gw1");
        assert_eq!(ch.gateway, "gw1");
        assert_eq!(ch.current_state(), ChannelState::Init);
        assert!(!ch.is_gathering());
        assert!(!ch.is_closed());
    }

    #[test]
    fn test_channel_start() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topology data");
        assert_eq!(ch.current_state(), ChannelState::Config);
        assert!(ch.channel.opened);

        // Should have enqueued a configuration message
        let outgoing = ch.channel.take_outgoing().unwrap();
        let xml = String::from_utf8(outgoing).unwrap();
        assert!(xml.contains("type=\"CFG\""));
        assert!(xml.contains("gateway=\"gw1\""));
    }

    #[test]
    fn test_channel_config_ack() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        assert_eq!(ch.state, ChannelState::Config);

        // Receive ACK
        let ack = Message::ack(0);
        let results = ch.recv(ack).unwrap();
        assert!(results.is_empty());
        assert_eq!(ch.state, ChannelState::Gather);
        assert!(ch.channel.setup);
    }

    #[test]
    fn test_channel_shell_command() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");

        // ACK the config
        ch.recv(Message::ack(0)).unwrap();

        // Send shell command
        ch.shell("node[1-4]", b"uname -a", 42);

        // Should have outgoing message
        let outgoing = ch.channel.take_outgoing().unwrap();
        // skip the cfg message
        let outgoing2 = ch.channel.take_outgoing().unwrap();
        let xml = String::from_utf8(outgoing2).unwrap();
        assert!(xml.contains("type=\"CTL\""));
        assert!(xml.contains("action=\"shell\""));
        assert!(xml.contains("target=\"node[1-4]\""));
    }

    #[test]
    fn test_channel_recv_stdout() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        // Simulate stdout message
        let mut stdout = Message::stdout("node[1-2]", 1);
        stdout.data_encode(b"hello world");
        let results = ch.recv(stdout).unwrap();

        assert_eq!(results.len(), 2);
        for r in &results {
            if let PropagationResult::StdOut { data, .. } = r {
                assert_eq!(data, b"hello world");
            } else {
                panic!("expected StdOut");
            }
        }
    }

    #[test]
    fn test_channel_recv_stderr() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let mut stderr = Message::stderr("node1", 1);
        stderr.data_encode(b"error output");
        let results = ch.recv(stderr).unwrap();

        assert_eq!(results.len(), 1);
        if let PropagationResult::StdErr { node, data } = &results[0] {
            assert_eq!(node, "node1");
            assert_eq!(data, b"error output");
        } else {
            panic!("expected StdErr");
        }
    }

    #[test]
    fn test_channel_recv_retcode() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let ret = Message::retcode("node[1-3]", 0, 1);
        let results = ch.recv(ret).unwrap();

        assert_eq!(results.len(), 3);
        for r in &results {
            if let PropagationResult::Retcode { rc, .. } = r {
                assert_eq!(*rc, 0);
            } else {
                panic!("expected Retcode");
            }
        }
    }

    #[test]
    fn test_channel_recv_timeout() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let tim = Message::timeout("node[5-6]", 1);
        let results = ch.recv(tim).unwrap();

        assert_eq!(results.len(), 2);
        for r in &results {
            assert!(matches!(r, PropagationResult::Timeout { .. }));
        }
    }

    #[test]
    fn test_channel_recv_routing() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let rtr = Message::routing("connect", "gw2", "node[10-20]", 1);
        let results = ch.recv(rtr).unwrap();

        assert_eq!(results.len(), 1);
        if let PropagationResult::Routing {
            event,
            gateway,
            targets,
        } = &results[0]
        {
            assert_eq!(event, "connect");
            assert_eq!(gateway, "gw2");
            assert_eq!(targets, "node[10-20]");
        } else {
            panic!("expected Routing");
        }
    }

    #[test]
    fn test_channel_recv_error() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let err = Message::error("something went wrong");
        let result = ch.recv(err);
        assert!(result.is_err());
    }

    #[test]
    fn test_channel_recv_end() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let end = Message::end();
        ch.recv(end).unwrap();
        assert!(ch.is_closed());
    }

    #[test]
    fn test_channel_ev_hup() {
        let mut ch = PropagationChannel::new("gw1");
        ch.ev_hup(1);
        assert_eq!(ch.last_rc, Some(1));
    }

    #[test]
    fn test_channel_ev_close() {
        let mut ch = PropagationChannel::new("gw1");
        ch.ev_close();
        assert!(ch.is_closed());
    }

    #[test]
    fn test_channel_send_queue() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        // Before ACK, messages should be queued
        ch.shell("node1", b"echo hi", 1);

        // Only the config message should be in outbox
        let _cfg = ch.channel.take_outgoing().unwrap();
        assert!(ch.channel.take_outgoing().is_none());

        // After ACK, queued message should be sent
        ch.recv(Message::ack(0)).unwrap();
        let ctl = ch.channel.take_outgoing().unwrap();
        let xml = String::from_utf8(ctl).unwrap();
        assert!(xml.contains("type=\"CTL\""));
    }

    #[test]
    fn test_channel_write_and_eof() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        ch.write("node1", b"input data", 1);
        ch.set_write_eof("node1", 1);

        let _ = ch.channel.take_outgoing(); // cfg
        let write_msg = ch.channel.take_outgoing().unwrap();
        let xml = String::from_utf8(write_msg).unwrap();
        assert!(xml.contains("action=\"write\""));

        let eof_msg = ch.channel.take_outgoing().unwrap();
        let xml2 = String::from_utf8(eof_msg).unwrap();
        assert!(xml2.contains("action=\"eof\""));
    }

    #[test]
    fn test_channel_results_accumulate() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        // Send multiple results
        let mut out1 = Message::stdout("node1", 1);
        out1.data_encode(b"line1");
        ch.recv(out1).unwrap();

        let mut out2 = Message::stdout("node2", 1);
        out2.data_encode(b"line2");
        ch.recv(out2).unwrap();

        let ret = Message::retcode("node[1-2]", 0, 1);
        ch.recv(ret).unwrap();

        // Should have accumulated 4 results (2 stdout + 2 retcode)
        assert_eq!(ch.results.len(), 4);
    }
}