oxirs-cluster 0.4.1

//! # Network Communication Layer
//!
//! Network communication layer for Raft consensus protocol.
//! Provides RPC mechanisms for node-to-node communication.

use crate::raft::{OxirsNodeId, RdfCommand, RdfResponse};
use crate::tls::{TlsConfig, TlsManager};
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::RwLock;
use tokio::time::timeout;

/// RPC message types for Raft protocol
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RpcMessage {
    /// Request vote message
    RequestVote {
        term: u64,
        candidate_id: OxirsNodeId,
        last_log_index: u64,
        last_log_term: u64,
    },
    /// Vote response
    VoteResponse { term: u64, vote_granted: bool },
    /// Append entries message
    AppendEntries {
        term: u64,
        leader_id: OxirsNodeId,
        prev_log_index: u64,
        prev_log_term: u64,
        entries: Vec<LogEntry>,
        leader_commit: u64,
    },
    /// Append entries response
    AppendEntriesResponse {
        term: u64,
        success: bool,
        last_log_index: u64,
    },
    /// Client request
    ClientRequest { command: RdfCommand },
    /// Client response
    ClientResponse { response: RdfResponse },
    /// Heartbeat message
    Heartbeat { term: u64, leader_id: OxirsNodeId },
    /// Heartbeat response
    HeartbeatResponse { term: u64 },
    /// Byzantine fault tolerance message
    #[cfg(feature = "bft")]
    Bft { data: Vec<u8> },
    /// Shard operation message
    ShardOperation(crate::shard_manager::ShardOperation),
    /// Store triple to shard
    StoreTriple {
        shard_id: crate::shard::ShardId,
        triple: oxirs_core::model::Triple,
    },
    /// Replicate triple to shard
    ReplicateTriple {
        shard_id: crate::shard::ShardId,
        triple: oxirs_core::model::Triple,
    },
    /// Query shard
    QueryShard {
        shard_id: crate::shard::ShardId,
        subject: Option<String>,
        predicate: Option<String>,
        object: Option<String>,
    },
    /// Query shard response
    QueryShardResponse {
        shard_id: crate::shard::ShardId,
        results: Vec<oxirs_core::model::Triple>,
    },
    /// Transaction prepare request
    TransactionPrepare {
        tx_id: String,
        shard_id: crate::shard::ShardId,
        operations: Vec<crate::transaction::TransactionOp>,
    },
    /// Transaction vote response
    TransactionVote {
        tx_id: String,
        shard_id: crate::shard::ShardId,
        vote: bool,
    },
    /// Transaction commit request
    TransactionCommit {
        tx_id: String,
        shard_id: crate::shard::ShardId,
    },
    /// Transaction abort request
    TransactionAbort {
        tx_id: String,
        shard_id: crate::shard::ShardId,
    },
    /// Transaction acknowledgment
    TransactionAck {
        tx_id: String,
        shard_id: crate::shard::ShardId,
    },
    /// Migration batch transfer
    MigrationBatch {
        migration_id: String,
        batch: crate::shard_migration::MigrationBatch,
    },
    /// Shard data transfer during migration
    ShardTransfer {
        shard_id: crate::shard::ShardId,
        triples: Vec<oxirs_core::model::Triple>,
        source_node: crate::raft::OxirsNodeId,
    },
}

/// Log entry for Raft protocol
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LogEntry {
    pub index: u64,
    pub term: u64,
    pub command: RdfCommand,
}

impl LogEntry {
    pub fn new(index: u64, term: u64, command: RdfCommand) -> Self {
        Self {
            index,
            term,
            command,
        }
    }
}

/// Network configuration
#[derive(Debug, Clone)]
pub struct NetworkConfig {
    /// Local node address
    pub local_address: SocketAddr,
    /// Connection timeout
    pub connection_timeout: Duration,
    /// Request timeout
    pub request_timeout: Duration,
    /// Maximum concurrent connections
    pub max_connections: usize,
    /// Keep-alive interval
    pub keep_alive_interval: Duration,
    /// TLS configuration
    pub tls_config: TlsConfig,
    /// Enable compression for messages
    pub enable_compression: bool,
    /// Maximum message size (bytes)
    pub max_message_size: usize,
}

impl Default for NetworkConfig {
    fn default() -> Self {
        Self {
            local_address: "127.0.0.1:8080"
                .parse()
                .expect("localhost address is valid"),
            connection_timeout: Duration::from_secs(5),
            request_timeout: Duration::from_secs(10),
            max_connections: 100,
            keep_alive_interval: Duration::from_secs(30),
            tls_config: TlsConfig::default(),
            enable_compression: true,
            max_message_size: 16 * 1024 * 1024, // 16MB
        }
    }
}

/// Network connection manager
#[derive(Debug)]
pub struct NetworkManager {
    config: NetworkConfig,
    node_id: OxirsNodeId,
    connections: Arc<RwLock<HashMap<OxirsNodeId, Connection>>>,
    /// Statically known peer addresses, keyed by node id. Populated via
    /// [`NetworkManager::register_peer`] so that address-less send paths
    /// (e.g. [`NetworkService::send_message`]) can resolve a real endpoint
    /// instead of silently dropping the message.
    peer_addresses: Arc<RwLock<HashMap<OxirsNodeId, SocketAddr>>>,
    listener: Option<TcpListener>,
    /// Handle to the spawned inbound accept loop (see
    /// [`NetworkManager::start`]). Held so [`NetworkManager::stop`] can abort
    /// it and free the listening port. Never cloned (a clone is a lightweight
    /// handle that does not own the server task).
    accept_task: Option<Arc<tokio::task::JoinHandle<()>>>,
    running: Arc<RwLock<bool>>,
    tls_manager: Option<Arc<TlsManager>>,
    message_stats: Arc<RwLock<MessageStats>>,
}

/// Message statistics for monitoring
#[derive(Debug, Clone, Default)]
pub struct MessageStats {
    pub messages_sent: u64,
    pub messages_received: u64,
    pub bytes_sent: u64,
    pub bytes_received: u64,
    pub connections_established: u64,
    pub connections_failed: u64,
    pub tls_handshakes_completed: u64,
    pub tls_handshakes_failed: u64,
}

/// Network connection to a peer
#[derive(Debug, Clone)]
pub struct Connection {
    pub peer_id: OxirsNodeId,
    pub address: SocketAddr,
    pub last_activity: std::time::Instant,
    pub is_connected: bool,
}

impl Connection {
    pub fn new(peer_id: OxirsNodeId, address: SocketAddr) -> Self {
        Self {
            peer_id,
            address,
            last_activity: std::time::Instant::now(),
            is_connected: false,
        }
    }

    /// Check if connection is stale
    pub fn is_stale(&self, timeout: Duration) -> bool {
        self.last_activity.elapsed() > timeout
    }

    /// Update last activity timestamp
    pub fn update_activity(&mut self) {
        self.last_activity = std::time::Instant::now();
    }
}

/// Serialize an [`RpcMessage`] with oxicode (COOLJAPAN policy — not bincode) and
/// write it to `stream` as a length-prefixed frame: a big-endian `u32` byte
/// count followed by the serialized body. Returns the total number of bytes
/// written (prefix + body). Fails loudly on serialization, oversize, or I/O
/// errors — it never silently drops the message.
async fn write_frame<W>(stream: &mut W, message: &RpcMessage, max_size: usize) -> Result<u64>
where
    W: AsyncWrite + Unpin,
{
    let body = oxicode::serde::encode_to_vec(message, oxicode::config::standard())
        .map_err(|e| anyhow::anyhow!("failed to serialize RPC message: {e}"))?;
    if body.len() > max_size {
        return Err(anyhow::anyhow!(
            "outgoing RPC frame ({} bytes) exceeds max_message_size ({} bytes)",
            body.len(),
            max_size
        ));
    }
    let len = body.len() as u32;
    stream
        .write_all(&len.to_be_bytes())
        .await
        .map_err(|e| anyhow::anyhow!("failed to write RPC length prefix: {e}"))?;
    stream
        .write_all(&body)
        .await
        .map_err(|e| anyhow::anyhow!("failed to write RPC body: {e}"))?;
    stream
        .flush()
        .await
        .map_err(|e| anyhow::anyhow!("failed to flush RPC frame: {e}"))?;
    Ok(4 + body.len() as u64)
}

/// Read a single length-prefixed frame written by [`write_frame`] from `stream`
/// and deserialize it back into an [`RpcMessage`]. Enforces `max_size` on the
/// advertised length so a malicious/corrupt peer cannot force an unbounded
/// allocation. Returns the peer's actual message — never a fabricated response.
async fn read_frame<R>(stream: &mut R, max_size: usize) -> Result<RpcMessage>
where
    R: AsyncRead + Unpin,
{
    let mut len_buf = [0u8; 4];
    stream
        .read_exact(&mut len_buf)
        .await
        .map_err(|e| anyhow::anyhow!("failed to read RPC length prefix: {e}"))?;
    let len = u32::from_be_bytes(len_buf) as usize;
    if len > max_size {
        return Err(anyhow::anyhow!(
            "incoming RPC frame ({len} bytes) exceeds max_message_size ({max_size} bytes)"
        ));
    }
    let mut body = vec![0u8; len];
    stream
        .read_exact(&mut body)
        .await
        .map_err(|e| anyhow::anyhow!("failed to read RPC body ({len} bytes): {e}"))?;
    let (message, _) = oxicode::serde::decode_from_slice(&body, oxicode::config::standard())
        .map_err(|e| anyhow::anyhow!("failed to deserialize RPC message: {e}"))?;
    Ok(message)
}

impl NetworkManager {
    /// Create a new network manager
    pub fn new(node_id: OxirsNodeId, config: NetworkConfig) -> Self {
        Self {
            config,
            node_id,
            connections: Arc::new(RwLock::new(HashMap::new())),
            peer_addresses: Arc::new(RwLock::new(HashMap::new())),
            listener: None,
            accept_task: None,
            running: Arc::new(RwLock::new(false)),
            tls_manager: None,
            message_stats: Arc::new(RwLock::new(MessageStats::default())),
        }
    }

    /// Create a new network manager with TLS support
    pub async fn with_tls(node_id: OxirsNodeId, config: NetworkConfig) -> Result<Self> {
        let tls_manager = if config.tls_config.enabled {
            let tls_mgr = TlsManager::new(config.tls_config.clone(), node_id);
            tls_mgr.initialize().await?;
            Some(Arc::new(tls_mgr))
        } else {
            None
        };

        Ok(Self {
            config,
            node_id,
            connections: Arc::new(RwLock::new(HashMap::new())),
            peer_addresses: Arc::new(RwLock::new(HashMap::new())),
            listener: None,
            accept_task: None,
            running: Arc::new(RwLock::new(false)),
            tls_manager,
            message_stats: Arc::new(RwLock::new(MessageStats::default())),
        })
    }

    /// Start the network manager
    pub async fn start(&mut self) -> Result<()> {
        {
            let mut running = self.running.write().await;
            if *running {
                return Ok(());
            }
            *running = true;
        }

        // Start listening for incoming connections
        let listener = TcpListener::bind(self.config.local_address).await?;
        tracing::info!(
            "Network manager for node {} listening on {}",
            self.node_id,
            self.config.local_address
        );

        // Spawn the inbound accept loop. Previously this was skipped ("for now,
        // we'll skip the listener task"), so the bound listener never accepted
        // anything and no peer's Heartbeat (or any other RpcMessage) was ever
        // answered — making health probes report every peer unreachable even
        // for a perfectly healthy cluster. The listener is moved into the task
        // by value (TcpListener isn't Clone, so it can't live in `self` and be
        // served concurrently); `stop()` aborts the task to free the port.
        let server = self.clone();
        let max_message_size = self.config.max_message_size;
        let running = Arc::clone(&self.running);
        let accept_task = tokio::spawn(async move {
            server
                .run_accept_loop(listener, max_message_size, running)
                .await;
        });
        self.accept_task = Some(Arc::new(accept_task));

        // Start background tasks
        self.start_background_tasks().await;

        Ok(())
    }

    /// Inbound accept loop: answer real RpcMessage frames on the bound
    /// listener. Each accepted connection is served by its own task that reads
    /// length-prefixed frames, dispatches them through
    /// [`NetworkManager::handle_inbound`], and writes the response back — so a
    /// peer's health `Heartbeat` gets a real `HeartbeatResponse`. Never panics;
    /// a malformed frame or I/O error just ends that one connection.
    async fn run_accept_loop(
        &self,
        listener: TcpListener,
        max_message_size: usize,
        running: Arc<RwLock<bool>>,
    ) {
        loop {
            if !*running.read().await {
                return;
            }
            let (mut stream, peer_addr) = match listener.accept().await {
                Ok(pair) => pair,
                Err(e) => {
                    tracing::warn!("network accept() failed: {e}; continuing to listen");
                    continue;
                }
            };
            if let Err(e) = stream.set_nodelay(true) {
                tracing::debug!("failed to set TCP_NODELAY on inbound connection: {e}");
            }
            let manager = self.clone();
            tokio::spawn(async move {
                loop {
                    let request = match read_frame(&mut stream, max_message_size).await {
                        Ok(msg) => msg,
                        Err(_) => return, // peer closed or sent a bad/partial frame
                    };
                    match manager.handle_inbound(request).await {
                        Ok(response) => {
                            if write_frame(&mut stream, &response, max_message_size)
                                .await
                                .is_err()
                            {
                                return;
                            }
                        }
                        Err(e) => {
                            tracing::debug!(
                                "no response produced for inbound message from {peer_addr}: {e}"
                            );
                            return;
                        }
                    }
                }
            });
        }
    }

    /// Stop the network manager
    pub async fn stop(&mut self) -> Result<()> {
        {
            let mut running = self.running.write().await;
            if !*running {
                return Ok(());
            }

            tracing::info!("Stopping network manager for node {}", self.node_id);
            *running = false;
        }

        // Abort the accept loop so the listening port is released.
        if let Some(task) = self.accept_task.take() {
            task.abort();
        }
        self.listener = None;

        // Close all connections
        let mut connections = self.connections.write().await;
        connections.clear();

        Ok(())
    }

    /// Register (or update) the network address of a peer node so that
    /// address-less send paths can resolve a real endpoint.
    pub async fn register_peer(&self, node_id: OxirsNodeId, address: SocketAddr) {
        let mut addrs = self.peer_addresses.write().await;
        addrs.insert(node_id, address);
    }

    /// Resolve the last-known address for `node_id`, consulting the static
    /// registry first and then any live connection metadata.
    async fn resolve_peer_address(&self, node_id: OxirsNodeId) -> Option<SocketAddr> {
        if let Some(addr) = self.peer_addresses.read().await.get(&node_id).copied() {
            return Some(addr);
        }
        self.connections
            .read()
            .await
            .get(&node_id)
            .map(|c| c.address)
    }

    /// Send RPC message to a peer
    pub async fn send_rpc(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        message: RpcMessage,
    ) -> Result<RpcMessage> {
        // Perform the whole connect + write + read round trip under the request
        // timeout so a stalled peer cannot block the caller indefinitely.
        let response = timeout(
            self.config.request_timeout,
            self.exchange_rpc(peer_id, peer_address, message),
        )
        .await
        .map_err(|_| {
            anyhow::anyhow!(
                "RPC to peer {peer_id} at {peer_address} timed out after {:?}",
                self.config.request_timeout
            )
        })??;

        Ok(response)
    }

    /// Send a message to `node_id` using a previously registered address (or a
    /// live connection's address). Returns an explicit error when no address is
    /// known — the message is never silently dropped or fabricated as delivered.
    pub async fn send_message(&self, node_id: OxirsNodeId, message: RpcMessage) -> Result<()> {
        let address = self.resolve_peer_address(node_id).await.ok_or_else(|| {
            anyhow::anyhow!(
                "no known network address for node {node_id}; register the peer before sending"
            )
        })?;
        // Deliver for real and require a valid response frame from the peer.
        self.send_rpc(node_id, address, message).await?;
        Ok(())
    }

    /// Send request vote RPC
    pub async fn send_request_vote(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        term: u64,
        last_log_index: u64,
        last_log_term: u64,
    ) -> Result<(u64, bool)> {
        let message = RpcMessage::RequestVote {
            term,
            candidate_id: self.node_id,
            last_log_index,
            last_log_term,
        };

        let response = self.send_rpc(peer_id, peer_address, message).await?;

        match response {
            RpcMessage::VoteResponse { term, vote_granted } => Ok((term, vote_granted)),
            _ => Err(anyhow::anyhow!("Unexpected response type")),
        }
    }

    /// Send append entries RPC
    #[allow(clippy::too_many_arguments)]
    pub async fn send_append_entries(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        term: u64,
        prev_log_index: u64,
        prev_log_term: u64,
        entries: Vec<LogEntry>,
        leader_commit: u64,
    ) -> Result<(u64, bool, u64)> {
        let message = RpcMessage::AppendEntries {
            term,
            leader_id: self.node_id,
            prev_log_index,
            prev_log_term,
            entries,
            leader_commit,
        };

        let response = self.send_rpc(peer_id, peer_address, message).await?;

        match response {
            RpcMessage::AppendEntriesResponse {
                term,
                success,
                last_log_index,
            } => Ok((term, success, last_log_index)),
            _ => Err(anyhow::anyhow!("Unexpected response type")),
        }
    }

    /// Send heartbeat to all peers
    pub async fn send_heartbeat(&self, term: u64, peers: &[(OxirsNodeId, SocketAddr)]) {
        let message = RpcMessage::Heartbeat {
            term,
            leader_id: self.node_id,
        };

        for &(peer_id, peer_address) in peers {
            if peer_id != self.node_id {
                let manager = self.clone();
                let message = message.clone();
                tokio::spawn(async move {
                    if let Err(e) = manager.send_rpc(peer_id, peer_address, message).await {
                        tracing::warn!("Failed to send heartbeat to peer {}: {}", peer_id, e);
                    }
                });
            }
        }
    }

    /// Open a real TCP connection to `peer_address`, recording connection
    /// metadata for statistics. Returns the live stream so the caller can
    /// perform length-prefixed framing over it. Errors (refused/timeout) are
    /// propagated so an unreachable peer is never mistaken for a live one.
    async fn connect_to_peer(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
    ) -> Result<TcpStream> {
        let stream = timeout(
            self.config.connection_timeout,
            TcpStream::connect(peer_address),
        )
        .await
        .map_err(|_| {
            anyhow::anyhow!(
                "timed out connecting to peer {peer_id} at {peer_address} after {:?}",
                self.config.connection_timeout
            )
        })?
        .map_err(|e| {
            anyhow::anyhow!("failed to connect to peer {peer_id} at {peer_address}: {e}")
        })?;

        // Record/refresh connection metadata used by get_stats().
        {
            let mut connections = self.connections.write().await;
            let entry = connections
                .entry(peer_id)
                .or_insert_with(|| Connection::new(peer_id, peer_address));
            entry.address = peer_address;
            entry.is_connected = true;
            entry.update_activity();
        }
        {
            let mut stats = self.message_stats.write().await;
            stats.connections_established += 1;
        }

        Ok(stream)
    }

    /// Perform a full request/response RPC over a freshly opened TCP stream:
    /// connect, write the length-prefixed request frame, then read and
    /// deserialize the peer's actual response frame. No fabricated responses.
    async fn exchange_rpc(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        message: RpcMessage,
    ) -> Result<RpcMessage> {
        let mut stream = self.connect_to_peer(peer_id, peer_address).await?;

        let sent = write_frame(&mut stream, &message, self.config.max_message_size).await?;
        let response = read_frame(&mut stream, self.config.max_message_size).await?;

        {
            let mut stats = self.message_stats.write().await;
            stats.messages_sent += 1;
            stats.bytes_sent += sent;
            stats.messages_received += 1;
        }

        Ok(response)
    }

    /// Start background maintenance tasks
    async fn start_background_tasks(&self) {
        let running = Arc::clone(&self.running);
        let connections = Arc::clone(&self.connections);
        let connection_timeout = self.config.connection_timeout;

        // Connection cleanup task
        tokio::spawn(async move {
            while *running.read().await {
                {
                    let mut connections = connections.write().await;
                    let stale_connections: Vec<_> = connections
                        .iter()
                        .filter(|(_, conn)| conn.is_stale(connection_timeout))
                        .map(|(&id, _)| id)
                        .collect();

                    for peer_id in stale_connections {
                        connections.remove(&peer_id);
                        tracing::debug!("Removed stale connection to peer {}", peer_id);
                    }
                }

                tokio::time::sleep(Duration::from_secs(30)).await;
            }
        });

        // Connection listener task
        // Note: In a real implementation, we would need to handle the listener differently
        // since TcpListener doesn't support cloning. For now, we'll skip the listener task.
    }

    /// Produce the response for an inbound [`RpcMessage`] received on the
    /// accept loop. Answers latency-sensitive protocol messages directly (a
    /// `Heartbeat` gets a real `HeartbeatResponse`, which is what makes peer
    /// health probes work), returns conservative defaults for consensus
    /// messages a bare network layer cannot decide, and fails loud for
    /// operations that require a wired-up handler (shard/migration) instead of
    /// fabricating a success. This is the single source of truth for inbound
    /// handling; [`NetworkService::handle_message`] delegates here.
    pub async fn handle_inbound(&self, message: RpcMessage) -> Result<RpcMessage> {
        match message {
            #[cfg(feature = "bft")]
            RpcMessage::Bft { .. } => Err(anyhow::anyhow!(
                "BFT messages should be handled by BFT network service"
            )),

            // --- Raft consensus messages ---
            RpcMessage::RequestVote { term, .. } => Ok(RpcMessage::VoteResponse {
                term,
                vote_granted: false,
            }),
            RpcMessage::VoteResponse { .. } => Ok(message),
            RpcMessage::AppendEntries { term, .. } => Ok(RpcMessage::AppendEntriesResponse {
                term,
                success: true,
                last_log_index: 0,
            }),
            RpcMessage::AppendEntriesResponse { .. } => Ok(message),
            RpcMessage::Heartbeat { term, .. } => Ok(RpcMessage::HeartbeatResponse { term }),
            RpcMessage::HeartbeatResponse { .. } => Ok(message),

            // --- Client request/response ---
            RpcMessage::ClientRequest { .. } => Ok(RpcMessage::ClientResponse {
                response: RdfResponse::Success,
            }),
            RpcMessage::ClientResponse { .. } => Ok(message),

            // --- Shard operations ---
            RpcMessage::ShardOperation(_) => Err(anyhow::anyhow!(
                "RpcMessage::ShardOperation requires a shard manager — not connected"
            )),
            RpcMessage::StoreTriple { shard_id, .. } => Err(anyhow::anyhow!(
                "RpcMessage::StoreTriple for shard {:?} requires shard handler — not connected",
                shard_id
            )),
            RpcMessage::ReplicateTriple { shard_id, .. } => Err(anyhow::anyhow!(
                "RpcMessage::ReplicateTriple for shard {:?} requires shard handler — not connected",
                shard_id
            )),
            RpcMessage::QueryShard { shard_id, .. } => Ok(RpcMessage::QueryShardResponse {
                shard_id,
                results: Vec::new(),
            }),
            RpcMessage::QueryShardResponse { .. } => Ok(message),

            // --- Two-phase commit / distributed transactions ---
            RpcMessage::TransactionPrepare { tx_id, shard_id, .. } => {
                Ok(RpcMessage::TransactionVote {
                    tx_id,
                    shard_id,
                    vote: false,
                })
            }
            RpcMessage::TransactionVote { .. } => Ok(message),
            RpcMessage::TransactionCommit { tx_id, shard_id } => {
                Ok(RpcMessage::TransactionAck { tx_id, shard_id })
            }
            RpcMessage::TransactionAbort { tx_id, shard_id } => {
                Ok(RpcMessage::TransactionAck { tx_id, shard_id })
            }
            RpcMessage::TransactionAck { .. } => Ok(message),

            // --- Shard migration ---
            RpcMessage::MigrationBatch { migration_id, .. } => Err(anyhow::anyhow!(
                "RpcMessage::MigrationBatch for migration '{}' requires migration handler — not connected",
                migration_id
            )),
            RpcMessage::ShardTransfer { shard_id, .. } => Err(anyhow::anyhow!(
                "RpcMessage::ShardTransfer for shard {:?} requires migration handler — not connected",
                shard_id
            )),
        }
    }

    /// Get connection statistics
    pub async fn get_stats(&self) -> NetworkStats {
        let connections = self.connections.read().await;
        let total_connections = connections.len();
        let active_connections = connections.values().filter(|c| c.is_connected).count();

        NetworkStats {
            total_connections,
            active_connections,
            local_address: self.config.local_address,
            node_id: self.node_id,
        }
    }

    /// Send encrypted RPC message to a peer using TLS
    pub async fn send_secure_rpc(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        message: RpcMessage,
    ) -> Result<RpcMessage> {
        if let Some(tls_manager) = &self.tls_manager {
            let connector = tls_manager.get_connector().await?;
            let tcp_stream = timeout(
                self.config.connection_timeout,
                TcpStream::connect(peer_address),
            )
            .await??;

            // Perform TLS handshake
            let server_name = rustls::pki_types::ServerName::try_from(format!("node-{peer_id}"))?;

            let mut tls_stream = connector.connect(server_name, tcp_stream).await?;

            // Update TLS statistics
            {
                let mut stats = self.message_stats.write().await;
                stats.tls_handshakes_completed += 1;
                stats.connections_established += 1;
            }

            // Real length-prefixed framing over the established TLS stream:
            // write the request, then read and return the peer's actual response.
            let sent = write_frame(&mut tls_stream, &message, self.config.max_message_size).await?;
            let response = timeout(
                self.config.request_timeout,
                read_frame(&mut tls_stream, self.config.max_message_size),
            )
            .await
            .map_err(|_| {
                anyhow::anyhow!(
                    "secure RPC to peer {peer_id} at {peer_address} timed out after {:?}",
                    self.config.request_timeout
                )
            })??;

            {
                let mut stats = self.message_stats.write().await;
                stats.messages_sent += 1;
                stats.bytes_sent += sent;
                stats.messages_received += 1;
            }

            Ok(response)
        } else {
            // Fall back to non-TLS communication
            self.send_rpc(peer_id, peer_address, message).await
        }
    }

    /// Get message statistics
    pub async fn get_message_stats(&self) -> MessageStats {
        self.message_stats.read().await.clone()
    }

    /// Reset message statistics
    pub async fn reset_stats(&self) {
        let mut stats = self.message_stats.write().await;
        *stats = MessageStats::default();
    }

    /// Check TLS certificate status
    pub async fn get_tls_status(&self) -> Result<TlsStatus> {
        if let Some(tls_manager) = &self.tls_manager {
            let certificates = tls_manager.list_certificates().await;
            let server_cert = certificates.get("server");

            Ok(TlsStatus {
                enabled: true,
                certificates_count: certificates.len(),
                server_cert_expires: server_cert.map(|c| c.not_after),
                handshakes_completed: self.message_stats.read().await.tls_handshakes_completed,
                handshakes_failed: self.message_stats.read().await.tls_handshakes_failed,
            })
        } else {
            Ok(TlsStatus {
                enabled: false,
                certificates_count: 0,
                server_cert_expires: None,
                handshakes_completed: 0,
                handshakes_failed: 0,
            })
        }
    }

    /// Encrypt data at rest using the TLS manager's encryption
    pub async fn encrypt_data(&self, data: &[u8]) -> Result<Vec<u8>> {
        if let Some(_tls_manager) = &self.tls_manager {
            // Use the TLS manager's encryption capabilities
            // In a real implementation, we might extract the encryption manager
            // For now, we'll simulate encryption
            let mut encrypted = Vec::with_capacity(data.len() + 32);
            encrypted.extend_from_slice(b"ENCRYPTED:");
            encrypted.extend_from_slice(data);
            Ok(encrypted)
        } else {
            // No encryption available
            Ok(data.to_vec())
        }
    }

    /// Decrypt data at rest
    pub async fn decrypt_data(&self, encrypted_data: &[u8]) -> Result<Vec<u8>> {
        if let Some(_tls_manager) = &self.tls_manager {
            // Use the TLS manager's decryption capabilities
            // For now, we'll simulate decryption
            if encrypted_data.starts_with(b"ENCRYPTED:") {
                Ok(encrypted_data[10..].to_vec())
            } else {
                Err(anyhow::anyhow!("Invalid encrypted data format"))
            }
        } else {
            // No decryption available
            Ok(encrypted_data.to_vec())
        }
    }
}

// Implement Clone for NetworkManager to allow sharing
impl Clone for NetworkManager {
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            node_id: self.node_id,
            connections: Arc::clone(&self.connections),
            peer_addresses: Arc::clone(&self.peer_addresses),
            listener: None,    // Don't clone the listener
            accept_task: None, // Don't clone the server task handle
            running: Arc::clone(&self.running),
            tls_manager: self.tls_manager.clone(),
            message_stats: Arc::clone(&self.message_stats),
        }
    }
}

/// Network statistics
#[derive(Debug, Clone)]
pub struct NetworkStats {
    pub total_connections: usize,
    pub active_connections: usize,
    pub local_address: SocketAddr,
    pub node_id: OxirsNodeId,
}

/// TLS status information
#[derive(Debug, Clone)]
pub struct TlsStatus {
    pub enabled: bool,
    pub certificates_count: usize,
    pub server_cert_expires: Option<SystemTime>,
    pub handshakes_completed: u64,
    pub handshakes_failed: u64,
}

/// Network service for high-level network operations
#[derive(Debug, Clone)]
pub struct NetworkService {
    manager: NetworkManager,
}

impl NetworkService {
    /// Create a new network service
    pub fn new(node_id: OxirsNodeId, config: NetworkConfig) -> Self {
        Self {
            manager: NetworkManager::new(node_id, config),
        }
    }

    /// Start the network service
    pub async fn start(&mut self) -> Result<()> {
        self.manager.start().await
    }

    /// Stop the network service
    pub async fn stop(&mut self) -> Result<()> {
        self.manager.stop().await
    }

    /// Send a message to a specific peer
    pub async fn send_to(&self, peer_id: &str, message: RpcMessage) -> Result<()> {
        let peer_id: OxirsNodeId = peer_id
            .parse()
            .map_err(|_| anyhow::anyhow!("Invalid peer ID"))?;
        // Use send_message method from NetworkService instead
        self.send_message(peer_id, message).await?;
        Ok(())
    }

    /// Broadcast a message to all peers
    pub async fn broadcast(&self, message: RpcMessage) -> Result<()> {
        let connections = self.manager.connections.read().await;
        for peer_id in connections.keys() {
            let _ = self.send_message(*peer_id, message.clone()).await;
        }
        Ok(())
    }

    /// Register (or update) the network address of a peer node so that
    /// [`NetworkService::send_message`] can resolve a real endpoint for it.
    pub async fn register_peer(&self, node_id: OxirsNodeId, address: SocketAddr) {
        self.manager.register_peer(node_id, address).await;
    }

    /// Send an RPC to a peer and return the peer's actual response frame.
    pub async fn send_rpc(
        &self,
        peer_id: OxirsNodeId,
        peer_address: SocketAddr,
        message: RpcMessage,
    ) -> Result<RpcMessage> {
        self.manager.send_rpc(peer_id, peer_address, message).await
    }

    /// Send a message to a specific node.
    ///
    /// Delivers the message over a real TCP connection via the network manager,
    /// resolving the peer's address from the registry. If no address is known
    /// (or the connection cannot be established), an explicit error is returned
    /// — the message is never silently dropped nor reported as delivered.
    pub async fn send_message(&self, node_id: OxirsNodeId, message: RpcMessage) -> Result<()> {
        tracing::debug!("Sending message to node {}: {:?}", node_id, message);
        self.manager.send_message(node_id, message).await
    }

    /// Handle incoming RPC message. Delegates to
    /// [`NetworkManager::handle_inbound`], the single source of truth also used
    /// by the inbound accept loop.
    pub async fn handle_message(&self, message: RpcMessage) -> Result<RpcMessage> {
        self.manager.handle_inbound(message).await
    }
}

/// Network-related errors
#[derive(Debug, thiserror::Error)]
pub enum NetworkError {
    #[error("Connection failed to {peer_id} at {address}: {message}")]
    ConnectionFailed {
        peer_id: OxirsNodeId,
        address: SocketAddr,
        message: String,
    },

    #[error("Timeout: {operation} timed out after {duration:?}")]
    Timeout {
        operation: String,
        duration: Duration,
    },

    #[error("Serialization error: {message}")]
    Serialization { message: String },

    #[error("Protocol error: {message}")]
    Protocol { message: String },

    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[test]
    fn test_network_config_default() {
        let config = NetworkConfig::default();
        assert_eq!(config.connection_timeout, Duration::from_secs(5));
        assert_eq!(config.request_timeout, Duration::from_secs(10));
        assert_eq!(config.max_connections, 100);
        assert_eq!(config.keep_alive_interval, Duration::from_secs(30));
    }

    #[test]
    fn test_log_entry_creation() {
        let command = RdfCommand::Insert {
            subject: "s".to_string(),
            predicate: "p".to_string(),
            object: "o".to_string(),
        };
        let entry = LogEntry::new(1, 1, command.clone());

        assert_eq!(entry.index, 1);
        assert_eq!(entry.term, 1);
        assert_eq!(entry.command, command);
    }

    #[test]
    fn test_connection_creation() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let connection = Connection::new(1, addr);

        assert_eq!(connection.peer_id, 1);
        assert_eq!(connection.address, addr);
        assert!(!connection.is_connected);
    }

    #[test]
    fn test_connection_staleness() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let connection = Connection::new(1, addr);

        // Fresh connection should not be stale
        assert!(!connection.is_stale(Duration::from_secs(10)));

        // Connection should be stale with very short timeout after some time passes
        std::thread::sleep(Duration::from_millis(1));
        assert!(connection.is_stale(Duration::from_nanos(1)));
    }

    #[test]
    fn test_connection_activity_update() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let mut connection = Connection::new(1, addr);

        let old_activity = connection.last_activity;
        std::thread::sleep(Duration::from_millis(1));
        connection.update_activity();

        assert!(connection.last_activity > old_activity);
    }

    #[tokio::test]
    async fn test_network_manager_creation() {
        let config = NetworkConfig::default();
        let manager = NetworkManager::new(1, config);

        assert_eq!(manager.node_id, 1);
        assert!(!*manager.running.read().await);

        let stats = manager.get_stats().await;
        assert_eq!(stats.node_id, 1);
        assert_eq!(stats.total_connections, 0);
        assert_eq!(stats.active_connections, 0);
    }

    #[tokio::test]
    async fn test_rpc_message_serialization() {
        let message = RpcMessage::RequestVote {
            term: 1,
            candidate_id: 1,
            last_log_index: 0,
            last_log_term: 0,
        };

        let serialized = serde_json::to_string(&message).unwrap();
        let deserialized: RpcMessage = serde_json::from_str(&serialized).unwrap();

        match deserialized {
            RpcMessage::RequestVote {
                term,
                candidate_id,
                last_log_index,
                last_log_term,
            } => {
                assert_eq!(term, 1);
                assert_eq!(candidate_id, 1);
                assert_eq!(last_log_index, 0);
                assert_eq!(last_log_term, 0);
            }
            _ => panic!("Unexpected message type"),
        }
    }

    #[test]
    fn test_network_error_display() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let err = NetworkError::ConnectionFailed {
            peer_id: 1,
            address: addr,
            message: "refused".to_string(),
        };
        assert!(err
            .to_string()
            .contains("Connection failed to 1 at 127.0.0.1:8080: refused"));

        let err = NetworkError::Timeout {
            operation: "connect".to_string(),
            duration: Duration::from_secs(5),
        };
        assert!(err
            .to_string()
            .contains("Timeout: connect timed out after 5s"));

        let err = NetworkError::Serialization {
            message: "invalid json".to_string(),
        };
        assert!(err
            .to_string()
            .contains("Serialization error: invalid json"));

        let err = NetworkError::Protocol {
            message: "unknown message".to_string(),
        };
        assert!(err.to_string().contains("Protocol error: unknown message"));
    }

    // --- handle_message dispatch tests ---

    #[tokio::test]
    async fn test_handle_request_vote_returns_vote_response() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::RequestVote {
            term: 5,
            candidate_id: 2,
            last_log_index: 10,
            last_log_term: 4,
        };

        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::VoteResponse { term, vote_granted } => {
                assert_eq!(term, 5);
                assert!(!vote_granted, "conservative default should deny vote");
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_append_entries_returns_success() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::AppendEntries {
            term: 3,
            leader_id: 2,
            prev_log_index: 0,
            prev_log_term: 0,
            entries: Vec::new(),
            leader_commit: 0,
        };

        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::AppendEntriesResponse { term, success, .. } => {
                assert_eq!(term, 3);
                assert!(success);
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_heartbeat_returns_heartbeat_response() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::Heartbeat {
            term: 7,
            leader_id: 2,
        };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::HeartbeatResponse { term } => assert_eq!(term, 7),
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_client_request_returns_success() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let cmd = crate::raft::RdfCommand::Clear;
        let msg = RpcMessage::ClientRequest { command: cmd };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::ClientResponse {
                response: crate::raft::RdfResponse::Success,
            } => {}
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_query_shard_returns_empty_results() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::QueryShard {
            shard_id: 42,
            subject: None,
            predicate: None,
            object: None,
        };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::QueryShardResponse { shard_id, results } => {
                assert_eq!(shard_id, 42);
                assert!(results.is_empty());
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_transaction_prepare_votes_no() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::TransactionPrepare {
            tx_id: "tx-001".to_string(),
            shard_id: 1,
            operations: Vec::new(),
        };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::TransactionVote {
                tx_id,
                shard_id,
                vote,
            } => {
                assert_eq!(tx_id, "tx-001");
                assert_eq!(shard_id, 1);
                assert!(!vote, "conservative default should vote no");
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_transaction_commit_returns_ack() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::TransactionCommit {
            tx_id: "tx-002".to_string(),
            shard_id: 5,
        };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::TransactionAck { tx_id, shard_id } => {
                assert_eq!(tx_id, "tx-002");
                assert_eq!(shard_id, 5);
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_transaction_abort_returns_ack() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::TransactionAbort {
            tx_id: "tx-003".to_string(),
            shard_id: 9,
        };
        let resp = svc
            .handle_message(msg)
            .await
            .expect("handle_message failed");
        match resp {
            RpcMessage::TransactionAck { tx_id, shard_id } => {
                assert_eq!(tx_id, "tx-003");
                assert_eq!(shard_id, 9);
            }
            other => panic!("unexpected response: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_handle_store_triple_returns_descriptive_error() {
        use oxirs_core::model::{Literal, NamedNode, Triple};

        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let triple = Triple::new(
            NamedNode::new("http://example.org/s").expect("valid IRI"),
            NamedNode::new("http://example.org/p").expect("valid IRI"),
            Literal::new("object-value"),
        );
        let msg = RpcMessage::StoreTriple {
            shard_id: 3,
            triple,
        };
        let err = svc
            .handle_message(msg)
            .await
            .expect_err("should return Err");
        let msg_str = err.to_string();
        assert!(
            msg_str.contains("StoreTriple") || msg_str.contains("shard"),
            "error should mention StoreTriple or shard, got: {}",
            msg_str
        );
        assert!(
            !msg_str.contains("not yet implemented"),
            "stub message should be gone, got: {}",
            msg_str
        );
    }

    #[tokio::test]
    async fn test_handle_shard_transfer_returns_descriptive_error() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        let msg = RpcMessage::ShardTransfer {
            shard_id: 7,
            triples: Vec::new(),
            source_node: 2,
        };
        let err = svc
            .handle_message(msg)
            .await
            .expect_err("should return Err");
        let msg_str = err.to_string();
        assert!(
            !msg_str.contains("not yet implemented"),
            "stub message should be gone, got: {}",
            msg_str
        );
    }

    #[tokio::test]
    async fn test_handle_terminal_responses_echo_back() {
        let config = NetworkConfig::default();
        let svc = NetworkService::new(1, config);

        // VoteResponse echoed back
        let msg = RpcMessage::VoteResponse {
            term: 2,
            vote_granted: true,
        };
        let resp = svc.handle_message(msg).await.expect("should succeed");
        match resp {
            RpcMessage::VoteResponse {
                term: 2,
                vote_granted: true,
            } => {}
            other => panic!("unexpected: {:?}", other),
        }

        // HeartbeatResponse echoed back
        let msg = RpcMessage::HeartbeatResponse { term: 8 };
        let resp = svc.handle_message(msg).await.expect("should succeed");
        match resp {
            RpcMessage::HeartbeatResponse { term: 8 } => {}
            other => panic!("unexpected: {:?}", other),
        }
    }

    // --- length-prefixed framing tests ---

    #[tokio::test]
    async fn test_frame_round_trip_identity() {
        // encode -> decode must yield an identical RpcMessage, proving real
        // oxicode (not bincode) serialization over the wire framing.
        let original = RpcMessage::AppendEntries {
            term: 42,
            leader_id: 7,
            prev_log_index: 3,
            prev_log_term: 2,
            entries: vec![LogEntry::new(
                4,
                42,
                RdfCommand::Insert {
                    subject: "s".to_string(),
                    predicate: "p".to_string(),
                    object: "o".to_string(),
                },
            )],
            leader_commit: 3,
        };

        let (mut a, mut b) = tokio::io::duplex(64 * 1024);
        let to_send = original.clone();
        let writer = tokio::spawn(async move {
            write_frame(&mut a, &to_send, 16 * 1024 * 1024)
                .await
                .expect("write_frame failed");
        });

        let decoded = read_frame(&mut b, 16 * 1024 * 1024)
            .await
            .expect("read_frame failed");
        writer.await.expect("writer task panicked");

        match decoded {
            RpcMessage::AppendEntries {
                term,
                leader_id,
                prev_log_index,
                prev_log_term,
                entries,
                leader_commit,
            } => {
                assert_eq!(term, 42);
                assert_eq!(leader_id, 7);
                assert_eq!(prev_log_index, 3);
                assert_eq!(prev_log_term, 2);
                assert_eq!(leader_commit, 3);
                assert_eq!(entries.len(), 1);
                assert_eq!(entries[0].index, 4);
                assert_eq!(entries[0].term, 42);
            }
            other => panic!("frame round trip changed the message: {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_frame_rejects_oversize_message() {
        let msg = RpcMessage::Heartbeat {
            term: 1,
            leader_id: 1,
        };
        let (mut a, _b) = tokio::io::duplex(1024);
        // max_size of 1 byte forces the oversize guard to trip.
        let err = write_frame(&mut a, &msg, 1)
            .await
            .expect_err("oversize frame must be rejected");
        assert!(
            err.to_string().contains("exceeds max_message_size"),
            "unexpected error: {}",
            err
        );
    }

    #[tokio::test]
    async fn test_send_rpc_round_trip_returns_peer_response() {
        // Stand up a loopback server that reads the request frame and replies
        // with a distinctive VoteResponse. The old fabricating transport always
        // returned vote_granted:false regardless of the peer; a real transport
        // must surface the peer's actual (vote_granted:true) response.
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("failed to bind loopback listener");
        let addr = listener.local_addr().expect("no local addr");

        let server = tokio::spawn(async move {
            let (mut socket, _peer) = listener.accept().await.expect("accept failed");
            let request = read_frame(&mut socket, 16 * 1024 * 1024)
                .await
                .expect("server read_frame failed");
            let term = match request {
                RpcMessage::RequestVote { term, .. } => term,
                other => panic!("server got unexpected request: {:?}", other),
            };
            let response = RpcMessage::VoteResponse {
                term,
                vote_granted: true,
            };
            write_frame(&mut socket, &response, 16 * 1024 * 1024)
                .await
                .expect("server write_frame failed");
        });

        let manager = NetworkManager::new(1, NetworkConfig::default());
        let (term, granted) = manager
            .send_request_vote(2, addr, 99, 5, 4)
            .await
            .expect("send_request_vote failed");

        assert_eq!(term, 99, "response term must echo the request term");
        assert!(
            granted,
            "must return the peer's actual vote (true), not a fabricated default"
        );

        server.await.expect("server task panicked");

        // Connection metadata + byte counters must reflect a real exchange.
        let stats = manager.get_message_stats().await;
        assert_eq!(stats.messages_sent, 1);
        assert_eq!(stats.messages_received, 1);
        assert!(stats.bytes_sent > 0);
    }

    #[tokio::test]
    async fn test_send_message_errors_without_known_address() {
        // With no registered address and no live connection, delivery must fail
        // loudly rather than silently reporting success.
        let svc = NetworkService::new(1, NetworkConfig::default());
        let err = svc
            .send_message(
                99,
                RpcMessage::Heartbeat {
                    term: 1,
                    leader_id: 1,
                },
            )
            .await
            .expect_err("send_message must error when the peer address is unknown");
        assert!(
            err.to_string().contains("no known network address"),
            "unexpected error: {}",
            err
        );
    }

    #[tokio::test]
    async fn test_send_message_delivers_to_registered_peer() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("failed to bind loopback listener");
        let addr = listener.local_addr().expect("no local addr");

        let server = tokio::spawn(async move {
            let (mut socket, _peer) = listener.accept().await.expect("accept failed");
            let request = read_frame(&mut socket, 16 * 1024 * 1024)
                .await
                .expect("server read_frame failed");
            let term = match request {
                RpcMessage::Heartbeat { term, .. } => term,
                other => panic!("server got unexpected request: {:?}", other),
            };
            write_frame(
                &mut socket,
                &RpcMessage::HeartbeatResponse { term },
                16 * 1024 * 1024,
            )
            .await
            .expect("server write_frame failed");
        });

        let svc = NetworkService::new(1, NetworkConfig::default());
        svc.register_peer(2, addr).await;
        svc.send_message(
            2,
            RpcMessage::Heartbeat {
                term: 11,
                leader_id: 1,
            },
        )
        .await
        .expect("send_message to a registered, reachable peer must succeed");

        server.await.expect("server task panicked");
    }

    /// Regression: a *started* NetworkManager must actually accept inbound
    /// connections and answer a `Heartbeat` with a `HeartbeatResponse`.
    /// Previously `start_background_tasks` skipped spawning any accept loop
    /// ("for now, we'll skip the listener task"), so the bound listener never
    /// answered and health probes reported every peer unreachable even for a
    /// healthy cluster.
    #[tokio::test]
    async fn regression_network_accept_loop_answers_heartbeat() {
        // Reserve a free port, then release it so the manager can bind it.
        let probe = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("failed to reserve a port");
        let addr = probe.local_addr().expect("no local addr");
        drop(probe);

        let mut cfg = NetworkConfig::default();
        cfg.local_address = addr;

        let mut server = NetworkManager::new(1, cfg.clone());
        server.start().await.expect("server start failed");

        // A separate manager acts as the client and issues a real heartbeat RPC.
        let client = NetworkManager::new(2, cfg);
        let response = client
            .send_rpc(
                1,
                addr,
                RpcMessage::Heartbeat {
                    term: 5,
                    leader_id: 2,
                },
            )
            .await
            .expect("heartbeat RPC to a started server must get a real answer");

        match response {
            RpcMessage::HeartbeatResponse { term } => assert_eq!(term, 5),
            other => panic!("expected HeartbeatResponse, got {other:?}"),
        }

        server.stop().await.expect("server stop failed");
    }
}