consortium-crate 0.3.1

//! 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),
}

/// Errors raised by [`PropagationChannel::recv`].
#[derive(Error, Debug)]
pub enum RecvError {
    /// Protocol/message-level failure (invalid payload, parse errors, ...).
    #[error("{0}")]
    Message(#[from] MessageProcessingError),

    /// Gateway runtime error received after channel setup. Upstream raises
    /// `TopologyError("%s: %s" % (gateway, reason))` here
    /// (PropagationChannel.recv_ctl); the rendered text is kept identical,
    /// e.g. "gw1: bad news".
    #[error("{0}")]
    Gateway(String),
}

// ============================================================================
// 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, with the target nodes
    /// whose pickup event must fire once the message is actually sent.
    /// Mirrors upstream's `self._sendq` of (ctl, pickup_worker, pickup_nodes).
    send_queue: VecDeque<(Message, Vec<String>)>,
    /// Pickup notifications produced when a CTL actually goes on the wire
    /// (upstream #594: `ev_pickup` fires only when the command really left
    /// the local process). Drained by the caller (TreeWorker), which
    /// forwards them as `ev_pickup`. write()/eof CTLs never add entries.
    pub pickup_events: Vec<String>,
    /// 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,
    },
}

/// Split a byte buffer into lines, mirroring Python's `bytes.splitlines()`
/// for the ASCII line boundaries used by the tree protocol (\n, \r\n, \r).
/// Lines are returned without their terminator; a trailing partial line is
/// kept. Matches the oracle's per-line stderr/stdout reporting.
fn split_lines(buf: &[u8]) -> Vec<Vec<u8>> {
    let mut lines = Vec::new();
    let mut start = 0;
    let mut i = 0;
    while i < buf.len() {
        if buf[i] == b'\n' || buf[i] == b'\r' {
            lines.push(buf[start..i].to_vec());
            if buf[i] == b'\r' && i + 1 < buf.len() && buf[i + 1] == b'\n' {
                i += 1;
            }
            i += 1;
            start = i;
        } else {
            i += 1;
        }
    }
    if start < buf.len() {
        lines.push(buf[start..].to_vec());
    }
    lines
}

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(),
            pickup_events: Vec::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;
    }

    /// Actually push a CTL message on the wire.
    ///
    /// Mirrors upstream `PropagationChannel._send_ctl`: if `pickup_nodes` is
    /// non-empty, a pickup notification is recorded for every target node
    /// AFTER the send — so a pickup event only fires when the command really
    /// left the local process (#594). write()/set_write_eof() CTLs never
    /// carry pickup nodes.
    fn send_ctl(&mut self, ctl: Message, pickup_nodes: Vec<String>) {
        self.channel.send(&ctl);
        self.pickup_events.extend(pickup_nodes);
    }

    /// 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) {
        self.send_queued_pickup(msg, Vec::new());
    }

    /// Helper used to send a message, using the msg queue if needed, with
    /// optional pickup nodes (mirrors upstream `send_queued(ctl,
    /// pickup_worker=..., pickup_nodes=...)`).
    fn send_queued_pickup(&mut self, msg: Message, pickup_nodes: Vec<String>) {
        if self.channel.setup && self.send_queue.is_empty() {
            // send now if channel is setup and sendq empty
            self.send_ctl(msg, pickup_nodes);
        } else {
            self.send_queue.push_back((msg, pickup_nodes));
        }
    }

    /// Send one queued message, if any (mirrors upstream `send_dequeue`:
    /// pickup notifications fire here too, at actual send time).
    fn send_dequeue(&mut self) {
        if let Some((msg, pickup_nodes)) = self.send_queue.pop_front() {
            self.send_ctl(msg, pickup_nodes);
        }
    }

    /// Create and enqueue a shell command for execution through the channel.
    ///
    /// Mirrors upstream `PropagationChannel.shell`: the target nodes are
    /// registered for pickup notification, which fires only when the CTL is
    /// actually written on the wire (#594).
    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);
        // Pass nodes so the pickup notification fires only when the CTL
        // actually leaves the local process (upstream #594).
        let pickup_nodes = NodeSet::parse(nodes)
            .map(|ns| ns.iter().collect())
            .unwrap_or_else(|_| vec![nodes.to_string()]);
        self.send_queued_pickup(ctl, pickup_nodes);
    }

    /// 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);
    }

    /// Report `buf` as stderr of `nodeset`, as seen from the gateway.
    ///
    /// Mirrors upstream `PropagationChannel._report_stderr`: a trailing
    /// newline is appended before splitting so that an empty buffer still
    /// yields one (empty) line (#249), and each line is reported separately
    /// for every node of the nodeset.
    fn report_stderr(&self, nodeset: &NodeSet, buf: &[u8]) -> Vec<PropagationResult> {
        let mut terminated = buf.to_vec();
        terminated.push(b'\n');
        let mut results = Vec::new();
        for line in split_lines(&terminated) {
            for node in nodeset.iter() {
                results.push(PropagationResult::StdErr {
                    node,
                    data: line.clone(),
                });
            }
        }
        results
    }

    /// 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>, RecvError> {
        let mut new_results = Vec::new();

        // EndMessage: close the channel (upstream: got EndMessage; closing)
        if msg.ident() == "END" {
            self.state = ChannelState::Closed;
            self.channel.close();
            return Ok(new_results);
        }

        // Non-routed stderr (srcid == 0): error messages received when the
        // channel is not established yet, or gateway-related messages, are
        // reported as stderr of the given nodes in ANY state — upstream
        // handles this before the setup dispatch (90d3195-range).
        if msg.ident() == "SER" && msg.srcid() == Some(0) {
            let decoded = msg.data_decode()?;
            let nodeset = NodeSet::parse(msg.nodes().unwrap_or_default()).unwrap_or_default();
            let results = self.report_stderr(&nodeset, &decoded);
            self.results.extend(results.clone());
            return Ok(results);
        }

        // StartMessage: gateway greeting opens the channel
        if msg.ident() == "CHA" {
            self.channel.opened = true;
            return Ok(new_results);
        }

        match self.state {
            ChannelState::Config => {
                // Waiting for ACK to configuration
                match msg.ident() {
                    "ACK" => {
                        self.channel.setup = true;
                        self.state = ChannelState::Gather;
                        self.send_dequeue();
                    }
                    "ERR" => {
                        // Gateway error before setup (eg. invalid CFG
                        // payload): report the reason as gateway stderr
                        // instead of dropping it (upstream 90d3195).
                        if let Message::Error { reason, .. } = &msg {
                            let gw = NodeSet::parse(&self.gateway).unwrap_or_default();
                            let results = self.report_stderr(&gw, reason.as_bytes());
                            self.results.extend(results.clone());
                            return Ok(results);
                        }
                    }
                    _ => {}
                }
            }
            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" => {
                        // Routed stderr: reported per line, like upstream
                        // recv_ctl's `decoded + b'\n'` splitlines() loop.
                        if let Some(nodes_str) = msg.nodes() {
                            if let Ok(decoded) = msg.data_decode() {
                                let ns = NodeSet::parse(nodes_str).unwrap_or_default();
                                new_results.extend(self.report_stderr(&ns, &decoded));
                            }
                        }
                    }
                    "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" => {
                        // Tree runtime error after setup: fatal to the
                        // channel (upstream raises TopologyError(
                        // "%s: %s" % (gateway, reason))).
                        if let Message::Error { reason, .. } = &msg {
                            self.channel.close();
                            self.state = ChannelState::Closed;
                            return Err(RecvError::Gateway(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);
    }

    // -- Upstream 1.10.1 ports (90d3195, Propagation._send_ctl/_report_stderr)

    /// Mirror of TreeHeadChannelErrorTest::test_head_channel_gw_error_before_setup:
    /// an ErrorMessage received before the gateway ACK is reported as stderr
    /// of the gateway node (instead of being dropped).
    #[test]
    fn test_head_channel_gw_error_before_setup() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        assert_eq!(ch.state, ChannelState::Config);

        let err =
            Message::error("Message CFG has an invalid payload (unsupported pickle protocol: 5)");
        let results = ch.recv(err).unwrap();

        assert_eq!(results.len(), 1);
        match &results[0] {
            PropagationResult::StdErr { node, data } => {
                assert_eq!(node, "gw1");
                assert_eq!(
                    data,
                    b"Message CFG has an invalid payload (unsupported pickle protocol: 5)"
                );
            }
            other => panic!("expected StdErr, got {:?}", other),
        }
        // not a fatal error at this stage: channel stays in Config
        assert_eq!(ch.state, ChannelState::Config);
    }

    /// Mirror of TreeHeadChannelErrorTest::test_head_channel_gw_error_empty_reason:
    /// an empty reason is still reported as one (empty) stderr line (#249).
    #[test]
    fn test_head_channel_gw_error_empty_reason() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");

        let results = ch.recv(Message::error("")).unwrap();
        assert_eq!(results.len(), 1);
        match &results[0] {
            PropagationResult::StdErr { node, data } => {
                assert_eq!(node, "gw1");
                assert_eq!(data, b"");
            }
            other => panic!("expected StdErr, got {:?}", other),
        }
    }

    /// Mirror of TreeHeadChannelErrorTest::test_head_channel_gw_error_after_setup:
    /// an ErrorMessage after setup raises a TopologyError-style failure
    /// ("gw: reason") and produces NO stderr results.
    #[test]
    fn test_head_channel_gw_error_after_setup() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();
        assert!(ch.channel.setup);

        let result = ch.recv(Message::error("bad news"));
        match result.unwrap_err() {
            RecvError::Gateway(text) => assert_eq!(text, "gw1: bad news"),
            other => panic!("expected Gateway error, got {:?}", other),
        }
        // no stderr results were produced
        assert!(ch.results.is_empty());
        // fatal: channel is closed
        assert!(ch.channel.closed);
    }

    /// Mirror of TreeHeadChannelErrorTest::test_head_channel_err_invalid_tag
    /// (logic-level): a malformed tag on the wire becomes a synthetic SER
    /// message (srcid 0) which the channel reports as gateway stderr; the
    /// channel is closed (fatal).
    #[test]
    fn test_head_channel_err_invalid_tag_reported_as_stderr() {
        let mut ch = PropagationChannel::new("gw1");
        // gateway greeting opens the channel
        let msgs = ch.channel.feed("<channel version=\"1.10.1\">").unwrap();
        for m in msgs {
            ch.recv(m).unwrap();
        }
        assert!(ch.channel.opened);

        // invalid tag on the wire: parse fails, initiator side converts to
        // a synthetic stderr message (6f45bf5), then closes the channel
        let err = ch.channel.feed("<foo></foo>").unwrap_err();
        let synthetic = ch
            .channel
            .handle_channel_error("gw1", &err)
            .expect("initiator must produce a synthetic stderr message");
        assert!(ch.channel.closed);

        let results = ch.recv(synthetic).unwrap();
        assert_eq!(results.len(), 1);
        match &results[0] {
            PropagationResult::StdErr { node, data } => {
                assert_eq!(node, "gw1");
                assert_eq!(data, b"Invalid starting tag foo");
            }
            other => panic!("expected StdErr, got {:?}", other),
        }
    }

    /// Mirror of TreeHeadChannelErrorTest::test_head_channel_err_parse_error:
    /// malformed XML is reported as gateway stderr containing "not
    /// well-formed", then the channel closes.
    #[test]
    fn test_head_channel_err_parse_error_reported_as_stderr() {
        let mut ch = PropagationChannel::new("gw1");
        let msgs = ch.channel.feed("<channel version=\"1.10.1\">").unwrap();
        for m in msgs {
            ch.recv(m).unwrap();
        }

        let err = ch
            .channel
            .feed("<message type=\"ABC\"</message>")
            .unwrap_err();
        let synthetic = ch
            .channel
            .handle_channel_error("gw1", &err)
            .expect("initiator must produce a synthetic stderr message");
        assert!(ch.channel.closed);

        let results = ch.recv(synthetic).unwrap();
        assert_eq!(results.len(), 1);
        match &results[0] {
            PropagationResult::StdErr { node, data } => {
                assert_eq!(node, "gw1");
                assert!(
                    String::from_utf8_lossy(data).contains("not well-formed"),
                    "expected 'not well-formed' in {:?}",
                    data
                );
            }
            other => panic!("expected StdErr, got {:?}", other),
        }
    }

    /// Non-routed stderr (SER with srcid 0) is reported per line in ANY
    /// state — here in Gather (upstream: handled before setup dispatch).
    #[test]
    fn test_channel_recv_nonrouted_stderr_srcid0() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let mut ser = Message::stderr("node[1-2]", 0);
        ser.data_encode(b"line a\nline b");
        let results = ch.recv(ser).unwrap();

        // 2 lines x 2 nodes, per-line ordering (line outer, node inner)
        let collected: Vec<(&str, &[u8])> = results
            .iter()
            .map(|r| match r {
                PropagationResult::StdErr { node, data } => (node.as_str(), data.as_slice()),
                other => panic!("expected StdErr, got {:?}", other),
            })
            .collect();
        assert_eq!(
            collected,
            vec![
                ("node1", b"line a".as_slice()),
                ("node2", b"line a".as_slice()),
                ("node1", b"line b".as_slice()),
                ("node2", b"line b".as_slice()),
            ]
        );
    }

    /// Routed stderr is also reported per line (upstream recv_ctl's
    /// `decoded + b'\n'` splitlines loop), without trailing newline kept.
    #[test]
    fn test_channel_recv_routed_stderr_per_line() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        let mut ser = Message::stderr("node1", 7);
        ser.data_encode(b"first\nsecond\n");
        let results = ch.recv(ser).unwrap();

        // like upstream `decoded + b'\n'` + splitlines(): a payload already
        // ending with '\n' yields a trailing empty line too
        let datas: Vec<&[u8]> = results
            .iter()
            .map(|r| match r {
                PropagationResult::StdErr { data, .. } => data.as_slice(),
                other => panic!("expected StdErr, got {:?}", other),
            })
            .collect();
        assert_eq!(
            datas,
            vec![b"first".as_slice(), b"second".as_slice(), b"".as_slice()]
        );
    }

    /// _report_stderr semantics: an empty buffer still yields one empty
    /// line per node (upstream #249 — trailing newline before splitlines).
    #[test]
    fn test_report_stderr_empty_buffer_yields_empty_line() {
        let ch = PropagationChannel::new("gw1");
        let ns = NodeSet::parse("node1").unwrap();
        let results = ch.report_stderr(&ns, b"");
        assert_eq!(results.len(), 1);
        match &results[0] {
            PropagationResult::StdErr { node, data } => {
                assert_eq!(node, "node1");
                assert_eq!(data, b"");
            }
            other => panic!("expected StdErr, got {:?}", other),
        }
    }

    /// CRLF and CR line boundaries split like Python's bytes.splitlines().
    #[test]
    fn test_split_lines_boundaries() {
        assert_eq!(
            split_lines(b"a\r\nb\rc\n"),
            vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]
        );
        assert_eq!(split_lines(b"no-newline"), vec![b"no-newline".to_vec()]);
        assert_eq!(split_lines(b"\n"), vec![b"".to_vec()]);
        assert!(split_lines(b"").is_empty());
    }

    /// #594 pickup semantics: shell() registers pickup nodes, but the
    /// notification fires only when the CTL actually goes on the wire —
    /// immediately when setup, or later via send_dequeue after the ACK.
    #[test]
    fn test_send_ctl_pickup_fires_on_actual_send() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");

        // queued before setup: no pickup yet
        ch.shell("node[1-2]", b"echo hi", 1);
        assert!(ch.pickup_events.is_empty());

        // ACK dequeues and sends the CTL: pickup fires now
        ch.recv(Message::ack(0)).unwrap();
        assert_eq!(
            ch.pickup_events,
            vec!["node1".to_string(), "node2".to_string()]
        );

        // setup channel: shell sends immediately, pickup fires at once
        ch.shell("node3", b"echo hi", 1);
        assert_eq!(
            ch.pickup_events,
            vec![
                "node1".to_string(),
                "node2".to_string(),
                "node3".to_string()
            ]
        );
    }

    /// write() and set_write_eof() never carry pickup nodes (upstream:
    /// "write()/set_write_eof() don't carry pickup args").
    #[test]
    fn test_write_and_eof_no_pickup() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");
        ch.recv(Message::ack(0)).unwrap();

        ch.write("node1", b"input", 1);
        ch.set_write_eof("node1", 1);
        assert!(ch.pickup_events.is_empty());
    }

    /// The queue preserves FIFO order with pickup metadata attached
    /// (upstream appendleft/pop deque semantics).
    #[test]
    fn test_send_queue_fifo_with_pickup() {
        let mut ch = PropagationChannel::new("gw1");
        ch.start(b"topo");

        ch.shell("node1", b"cmd1", 1);
        ch.shell("node2", b"cmd2", 2);

        let _cfg = ch.channel.take_outgoing().unwrap();
        assert!(ch.channel.take_outgoing().is_none());

        ch.recv(Message::ack(0)).unwrap();
        let first = String::from_utf8(ch.channel.take_outgoing().unwrap()).unwrap();
        assert!(first.contains("target=\"node1\""));
        // only the first queued CTL was dequeued by the ACK
        assert_eq!(ch.pickup_events, vec!["node1".to_string()]);
    }
}