//! The connection to the node this server is bound to.
//!
//! `ConnectionManager` owns that one connection for the life of the server,
//! and `NodeConnection` is the connection itself.

use crate::error::{ConnectionError, ConnectionResult};
use crate::server::FormatterMode;
use eetf::{Atom, BigInteger, Term, Tuple};
use erl_dist::epmd::EpmdClient;
use erl_dist::handshake::{ClientSideHandshake, HandshakeStatus};
use erl_dist::message::{self, Message, Receiver, Sender};
use erl_dist::node::{Creation, LocalNode, NodeName, PeerNode};
use futures_io::{AsyncRead, AsyncWrite};
use std::collections::HashMap;
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::task::{Context, Poll};
use std::time::Instant;
use tokio::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite, ReadBuf};
use tokio::net::TcpStream;
use tokio::sync::{Mutex, RwLock, mpsc};
use tokio::task::JoinHandle;

/// A wrapper around a stream that implements Clone and futures_io traits.
/// Clone is implemented by wrapping in Arc, allowing shared access.
pub(crate) struct ClonableStream<S> {
    inner: Arc<Mutex<S>>,
}

impl<S> ClonableStream<S> {
    fn new(stream: S) -> Self {
        Self {
            inner: Arc::new(Mutex::new(stream)),
        }
    }
}

impl<S> Clone for ClonableStream<S> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<S: TokioAsyncRead + Unpin> AsyncRead for ClonableStream<S> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        let mut guard = match self.inner.try_lock() {
            Ok(guard) => guard,
            Err(_) => {
                cx.waker().wake_by_ref();
                return Poll::Pending;
            }
        };
        let stream = &mut *guard;
        let mut read_buf = ReadBuf::new(buf);
        match Pin::new(stream).poll_read(cx, &mut read_buf) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(read_buf.filled().len())),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<S: TokioAsyncWrite + Unpin> AsyncWrite for ClonableStream<S> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let mut guard = match self.inner.try_lock() {
            Ok(guard) => guard,
            Err(_) => {
                cx.waker().wake_by_ref();
                return Poll::Pending;
            }
        };
        let stream = &mut *guard;
        Pin::new(stream).poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut guard = match self.inner.try_lock() {
            Ok(guard) => guard,
            Err(_) => {
                cx.waker().wake_by_ref();
                return Poll::Pending;
            }
        };
        let stream = &mut *guard;
        Pin::new(stream).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut guard = match self.inner.try_lock() {
            Ok(guard) => guard,
            Err(_) => {
                cx.waker().wake_by_ref();
                return Poll::Pending;
            }
        };
        let stream = &mut *guard;
        Pin::new(stream).poll_shutdown(cx)
    }
}

/// The name this server presents to a node when nothing else is given.
pub(crate) const DEFAULT_LOCAL_NODE: &str = "beamdev@localhost";

/// Atom identifying our correlation tags in `$gen_call` requests.
const CALL_REF_TAG: &str = "$beamdev_call";

/// Correlation tag embedded in a `$gen_call` request and echoed back verbatim
/// by the peer in its `{Tag, Result}` reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct CallRef(u64);

impl CallRef {
    /// Allocates the next correlation tag.
    pub(crate) fn next() -> Self {
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        Self(COUNTER.fetch_add(1, Ordering::Relaxed))
    }

    /// Encodes this tag as the term to send as the `$gen_call` reply address.
    pub(crate) fn to_term(self) -> Term {
        Term::from(Tuple::from(vec![
            Term::from(Atom::from(CALL_REF_TAG)),
            Term::from(BigInteger::from(self.0)),
        ]))
    }

    /// Recovers a tag from a reply, which the peer re-encodes in its own
    /// canonical integer representation.
    fn from_term(term: &Term) -> Option<Self> {
        let Term::Tuple(tuple) = term else {
            return None;
        };
        let [Term::Atom(tag), id] = tuple.elements.as_slice() else {
            return None;
        };
        if tag.name != CALL_REF_TAG {
            return None;
        }
        match id {
            Term::FixInteger(id) => u64::try_from(id.value).ok().map(Self),
            Term::BigInteger(id) => u64::try_from(&id.value).ok().map(Self),
            _ => None,
        }
    }
}

/// Request sent to the background connection task for RPC calls.
#[derive(Debug)]
pub(crate) struct RpcRequest {
    /// The message to send.
    pub(crate) message: Message,
    /// Correlation tag the reply will carry.
    pub(crate) call_ref: CallRef,
    /// Channel to send the response back.
    pub(crate) response_tx: mpsc::Sender<RpcResponse>,
}

/// Response from an RPC call.
#[derive(Debug)]
pub(crate) enum RpcResponse {
    /// Successfully received a response term.
    Success(Term),
    /// The RPC failed with an error.
    Error(ConnectionError),
}

/// State of a node connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ConnectionState {
    /// Connection is being established.
    Connecting,
    /// Connection is active and ready for use.
    Connected,
    /// Connection has been disconnected.
    Disconnected,
}

impl std::fmt::Display for ConnectionState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ConnectionState::Connecting => write!(f, "connecting"),
            ConnectionState::Connected => write!(f, "connected"),
            ConnectionState::Disconnected => write!(f, "disconnected"),
        }
    }
}

/// Status information about a connected node.
#[derive(Debug, Clone)]
pub(crate) struct NodeStatus {
    /// The node name.
    pub(crate) name: String,
    /// Current connection state.
    pub(crate) state: ConnectionState,
    /// When the connection was established (if connected).
    pub(crate) connected_at: Option<Instant>,
}

/// A connection to a single Erlang node.
pub(crate) struct NodeConnection {
    /// The node name (e.g., "foo@localhost").
    pub(crate) node_name: String,
    /// Current connection state.
    state: Arc<RwLock<ConnectionState>>,
    /// When the connection was established.
    connected_at: Option<Instant>,
    /// Handle to the background task managing this connection.
    task_handle: Option<JoinHandle<()>>,
    /// Channel for sending RPC requests to the background task.
    request_tx: Option<mpsc::Sender<RpcRequest>>,
    /// Information about the peer node.
    peer_node: Option<PeerNode>,
}

impl std::fmt::Debug for NodeConnection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NodeConnection")
            .field("node_name", &self.node_name)
            .field("state", &self.state)
            .field("connected_at", &self.connected_at)
            .finish_non_exhaustive()
    }
}

impl NodeConnection {
    /// Creates a new NodeConnection in the Connecting state.
    fn new(node_name: String) -> Self {
        Self {
            node_name,
            state: Arc::new(RwLock::new(ConnectionState::Connecting)),
            connected_at: None,
            task_handle: None,
            request_tx: None,
            peer_node: None,
        }
    }

    /// Returns the current connection state.
    pub(crate) async fn state(&self) -> ConnectionState {
        *self.state.read().await
    }

    /// Returns the status of this connection.
    pub(crate) async fn status(&self) -> NodeStatus {
        NodeStatus {
            name: self.node_name.clone(),
            state: self.state().await,
            connected_at: self.connected_at,
        }
    }

    /// Returns a handle for issuing requests without holding the connection lock.
    pub(crate) fn request_sender(&self) -> ConnectionResult<RequestSender> {
        self.request_tx
            .clone()
            .map(|request_tx| RequestSender {
                node_name: self.node_name.clone(),
                request_tx,
            })
            .ok_or_else(|| ConnectionError::NotConnected {
                node: self.node_name.clone(),
            })
    }

    /// Gracefully shuts down this connection.
    async fn shutdown(&mut self) {
        self.request_tx.take();

        if let Some(handle) = self.task_handle.take() {
            let _ = handle.await;
        }

        *self.state.write().await = ConnectionState::Disconnected;
    }
}

/// Handle for issuing requests on an established connection.
#[derive(Debug, Clone)]
pub(crate) struct RequestSender {
    node_name: String,
    request_tx: mpsc::Sender<RpcRequest>,
}

impl RequestSender {
    /// Sends a request and awaits the reply carrying `call_ref`.
    pub(crate) async fn call(&self, message: Message, call_ref: CallRef) -> ConnectionResult<Term> {
        let (response_tx, mut response_rx) = mpsc::channel(1);
        let request = RpcRequest {
            message,
            call_ref,
            response_tx,
        };

        self.request_tx
            .send(request)
            .await
            .map_err(|_| ConnectionError::ConnectionLost {
                node: self.node_name.clone(),
                reason: "request channel closed".to_string(),
            })?;

        match response_rx.recv().await {
            Some(RpcResponse::Success(term)) => Ok(term),
            Some(RpcResponse::Error(err)) => Err(err),
            None => Err(ConnectionError::ConnectionLost {
                node: self.node_name.clone(),
                reason: "response channel closed".to_string(),
            }),
        }
    }
}

/// Owns the connection to the one node this server works against.
pub struct ConnectionManager {
    /// The node we connect to, named at startup.
    node: String,
    /// The connection to it, once it has been established.
    connection: RwLock<Option<Arc<Mutex<NodeConnection>>>>,
    /// Our local node name.
    local_node_name: String,
    /// How many times the eval runner has been loaded onto the node.
    runner_generation: Arc<Mutex<u64>>,
    /// The syntax terms are rendered in when they appear in error messages.
    formatter_mode: FormatterMode,
}

impl std::fmt::Debug for ConnectionManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConnectionManager")
            .field("node", &self.node)
            .field("local_node_name", &self.local_node_name)
            .finish_non_exhaustive()
    }
}

impl ConnectionManager {
    /// Creates a manager for `node`, connecting under `local_node_name`.
    ///
    /// Each server presents a distinct local name: a node rejects concurrent
    /// handshakes from the same peer name.
    pub fn new(local_node_name: String, node: String) -> Self {
        Self {
            node,
            connection: RwLock::new(None),
            local_node_name,
            runner_generation: Arc::new(Mutex::new(0)),
            formatter_mode: FormatterMode::default(),
        }
    }

    /// Renders terms in `mode` when they appear in error messages.
    pub(crate) fn with_formatter_mode(mut self, mode: FormatterMode) -> Self {
        self.formatter_mode = mode;
        self
    }

    /// The syntax terms are rendered in when they appear in error messages.
    pub(crate) fn formatter_mode(&self) -> FormatterMode {
        self.formatter_mode
    }

    /// The node this manager connects to.
    pub(crate) fn node(&self) -> &str {
        &self.node
    }

    /// How many times the eval runner has been loaded onto the node.
    ///
    /// Loading a module purges the version before it, killing every process
    /// still running that code -- so a redundant deploy kills the runner
    /// processes of calls already in flight. Callers racing to replace the same
    /// missing runner compare generations to see that someone else got there.
    pub(crate) fn runner_generation(&self) -> Arc<Mutex<u64>> {
        Arc::clone(&self.runner_generation)
    }

    /// The node name this server presents to the nodes it connects to.
    pub(crate) fn local_node_name(&self) -> &str {
        &self.local_node_name
    }

    /// Connects to the node, replacing any connection already established.
    ///
    /// # Errors
    /// Returns an error if the node is unreachable, authentication fails, or
    /// the handshake fails.
    pub async fn connect(&self, cookie: &str) -> ConnectionResult<()> {
        let node_name = &self.node;
        let mut node_conn = NodeConnection::new(node_name.clone());

        let (short_name, host) = parse_node_name(node_name)?;

        let epmd_addr = format!("{}:{}", host, erl_dist::epmd::DEFAULT_EPMD_PORT);
        let epmd_stream =
            TcpStream::connect(&epmd_addr)
                .await
                .map_err(|e| ConnectionError::NodeUnreachable {
                    node: node_name.clone(),
                    reason: format!("cannot connect to EPMD at {}: {}", epmd_addr, e),
                })?;

        let epmd_client = EpmdClient::new(ClonableStream::new(epmd_stream));
        let node_entry = epmd_client
            .get_node(&short_name)
            .await
            .map_err(|e| ConnectionError::NodeUnreachable {
                node: node_name.clone(),
                reason: format!("EPMD lookup failed: {}", e),
            })?
            .ok_or_else(|| ConnectionError::NodeUnreachable {
                node: node_name.clone(),
                reason: "node not registered with EPMD".to_string(),
            })?;

        let node_addr = format!("{}:{}", host, node_entry.port);
        let stream =
            TcpStream::connect(&node_addr)
                .await
                .map_err(|e| ConnectionError::NodeUnreachable {
                    node: node_name.clone(),
                    reason: format!("cannot connect to node at {}: {}", node_addr, e),
                })?;

        let clonable_stream = ClonableStream::new(stream);

        let local_node_name: NodeName =
            self.local_node_name
                .parse()
                .map_err(|e| ConnectionError::HandshakeFailed {
                    node: node_name.clone(),
                    reason: format!("invalid local node name: {}", e),
                })?;

        let local_node = LocalNode::new(local_node_name, Creation::random());

        let mut handshake = ClientSideHandshake::new(clonable_stream, local_node, cookie);

        let status = handshake
            .execute_send_name(erl_dist::LOWEST_DISTRIBUTION_PROTOCOL_VERSION)
            .await
            .map_err(|e| ConnectionError::HandshakeFailed {
                node: node_name.clone(),
                reason: format!("send_name failed: {}", e),
            })?;

        let do_continue = matches!(status, HandshakeStatus::Alive);

        let (stream, peer_node) = handshake.execute_rest(do_continue).await.map_err(|e| {
            let err_str = e.to_string().to_lowercase();
            if err_str.contains("not_allowed") || err_str.contains("authentication") {
                ConnectionError::AuthenticationFailed {
                    node: node_name.clone(),
                }
            } else {
                ConnectionError::HandshakeFailed {
                    node: node_name.clone(),
                    reason: format!("handshake failed: {}", e),
                }
            }
        })?;

        let (sender, receiver) = message::channel(stream, peer_node.flags);

        let (request_tx, request_rx) = mpsc::channel::<RpcRequest>(32);

        let state_clone = node_conn.state.clone();
        let node_name_clone = node_name.clone();
        let task_handle = tokio::spawn(async move {
            connection_task(sender, receiver, request_rx, state_clone, node_name_clone).await;
        });

        node_conn.peer_node = Some(peer_node);
        node_conn.task_handle = Some(task_handle);
        node_conn.request_tx = Some(request_tx);
        node_conn.connected_at = Some(Instant::now());
        *node_conn.state.write().await = ConnectionState::Connected;

        self.replace(node_conn).await;

        Ok(())
    }

    /// Installs `connection` as the current one, shutting down any predecessor.
    async fn replace(&self, connection: NodeConnection) {
        let previous = self
            .connection
            .write()
            .await
            .replace(Arc::new(Mutex::new(connection)));

        if let Some(previous) = previous {
            previous.lock().await.shutdown().await;
        }
    }

    /// Drops the connection, if there is one.
    pub async fn disconnect(&self) {
        let connection = self.connection.write().await.take();

        if let Some(connection) = connection {
            connection.lock().await.shutdown().await;
        }
    }

    /// The status of the connection, or `None` if we have never connected.
    pub(crate) async fn status(&self) -> Option<NodeStatus> {
        let connection = self.connection().await?;
        let status = connection.lock().await.status().await;
        Some(status)
    }

    /// The current connection, if there is one.
    pub(crate) async fn connection(&self) -> Option<Arc<Mutex<NodeConnection>>> {
        self.connection.read().await.clone()
    }

    /// Whether the connection is up.
    pub async fn is_connected(&self) -> bool {
        match self.connection().await {
            Some(connection) => connection.lock().await.state().await == ConnectionState::Connected,
            None => false,
        }
    }
}

/// The peer end of an in-memory connection, for tests.
#[cfg(test)]
pub(crate) type TestPeer = (
    Sender<ClonableStream<tokio::io::DuplexStream>>,
    Receiver<ClonableStream<tokio::io::DuplexStream>>,
);

#[cfg(test)]
impl ConnectionManager {
    /// A manager for `node`, under the default local name.
    pub(crate) fn for_test(node: &str) -> Self {
        Self::new(DEFAULT_LOCAL_NODE.to_string(), node.to_string())
    }

    /// Registers a connection backed by an in-memory stream and returns the
    /// peer's end of the message channel.
    pub(crate) async fn connect_in_memory(&self) -> TestPeer {
        let node_name = self.node.clone();
        let (local, remote) = tokio::io::duplex(1024 * 1024);
        let (sender, receiver) = message::channel(
            ClonableStream::new(local),
            erl_dist::DistributionFlags::default(),
        );
        let peer = message::channel(
            ClonableStream::new(remote),
            erl_dist::DistributionFlags::default(),
        );

        let (request_tx, request_rx) = mpsc::channel(32);
        let mut conn = NodeConnection::new(node_name.clone());
        let state = conn.state.clone();
        conn.task_handle = Some(tokio::spawn(connection_task(
            sender, receiver, request_rx, state, node_name,
        )));

        conn.request_tx = Some(request_tx);
        conn.connected_at = Some(Instant::now());
        *conn.state.write().await = ConnectionState::Connected;

        self.replace(conn).await;

        peer
    }
}

/// Parses a node name into (short_name, host).
fn parse_node_name(node_name: &str) -> ConnectionResult<(String, String)> {
    let parts: Vec<&str> = node_name.split('@').collect();
    if parts.len() != 2 {
        return Err(ConnectionError::HandshakeFailed {
            node: node_name.to_string(),
            reason: "invalid node name format, expected 'name@host'".to_string(),
        });
    }
    Ok((parts[0].to_string(), parts[1].to_string()))
}

/// Background task that manages a node connection.
async fn connection_task<T>(
    mut sender: Sender<T>,
    mut receiver: Receiver<T>,
    mut request_rx: mpsc::Receiver<RpcRequest>,
    state: Arc<RwLock<ConnectionState>>,
    node_name: String,
) where
    T: AsyncRead + AsyncWrite + Unpin + Clone,
{
    let mut tick_interval = tokio::time::interval(tokio::time::Duration::from_secs(30));
    let mut pending: HashMap<CallRef, mpsc::Sender<RpcResponse>> = HashMap::new();

    loop {
        tokio::select! {
            request = request_rx.recv() => {
                match request {
                    Some(rpc_request) => {
                        if let Err(e) = sender.send(rpc_request.message).await {
                            let err = ConnectionError::ConnectionLost {
                                node: node_name.clone(),
                                reason: format!("send failed: {}", e),
                            };
                            let _ = rpc_request.response_tx.send(RpcResponse::Error(err)).await;
                            break;
                        }
                        pending.insert(rpc_request.call_ref, rpc_request.response_tx);
                    }
                    None => {
                        break;
                    }
                }
            }

            msg = receiver.recv() => {
                match msg {
                    Ok(message) => {
                        match message {
                            Message::Tick => {
                                tracing::trace!("Received tick from {}", node_name);
                            }
                            _ => {
                                let Some(term) = extract_message_term(message) else {
                                    continue;
                                };

                                match split_reply(term) {
                                    Ok((call_ref, result)) => match pending.remove(&call_ref) {
                                        Some(tx) => {
                                            let _ = tx.send(RpcResponse::Success(result)).await;
                                        }
                                        None => {
                                            tracing::debug!("Received reply for unknown call {:?} from {}", call_ref, node_name);
                                        }
                                    },
                                    Err(_) => {
                                        tracing::debug!("Received unsolicited message from {}", node_name);
                                    }
                                }
                            }
                        }
                    }
                    Err(e) => {
                        tracing::error!("Error receiving from {}: {}", node_name, e);
                        for (_, tx) in pending.drain() {
                            let _ = tx.send(RpcResponse::Error(ConnectionError::ConnectionLost {
                                node: node_name.clone(),
                                reason: format!("receive error: {}", e),
                            })).await;
                        }
                        break;
                    }
                }
            }

            _ = tick_interval.tick() => {
                sweep_abandoned(&mut pending);
                if let Err(e) = sender.send(Message::Tick).await {
                    tracing::error!("Failed to send tick to {}: {}", node_name, e);
                    break;
                }
            }
        }
    }

    *state.write().await = ConnectionState::Disconnected;
    tracing::info!("Connection task for {} terminated", node_name);
}

/// Drop in-flight calls whose caller has gone away, so a reply that never
/// arrives cannot pin an entry for the life of the connection.
fn sweep_abandoned(pending: &mut HashMap<CallRef, mpsc::Sender<RpcResponse>>) {
    pending.retain(|_, tx| !tx.is_closed());
}

/// Extract the Term payload from a received Message.
fn extract_message_term(message: Message) -> Option<Term> {
    match message {
        Message::Send(s) => Some(s.message),
        Message::SendSender(s) => Some(s.message),
        Message::SendTt(s) => Some(s.message),
        Message::SendSenderTt(s) => Some(s.message),
        Message::RegSend(s) => Some(s.message),
        Message::RegSendTt(s) => Some(s.message),
        Message::AliasSend(s) => Some(s.message),
        Message::AliasSendTt(s) => Some(s.message),
        _ => None,
    }
}

/// Split a `$gen_call` reply tuple `{Tag, Result}` into its correlation tag and
/// result, handing the term back when it is not a reply.
fn split_reply(term: Term) -> Result<(CallRef, Term), Term> {
    let Term::Tuple(tuple) = term else {
        return Err(term);
    };
    if tuple.elements.len() != 2 {
        return Err(Term::Tuple(tuple));
    }
    let Some(call_ref) = CallRef::from_term(&tuple.elements[0]) else {
        return Err(Term::Tuple(tuple));
    };
    let mut elements = tuple.elements;
    let result = elements.remove(1);
    Ok((call_ref, result))
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use eetf::{Atom, FixInteger, Pid, Tuple};
    use erl_dist::DistributionFlags;
    use std::time::Duration;

    fn test_pid() -> Pid {
        Pid::new("beamdev@localhost", 0, 0, 0)
    }

    fn gen_call(call_ref: CallRef) -> Message {
        Message::reg_send(
            test_pid(),
            Atom::from("rex"),
            Term::from(Tuple::from(vec![
                Term::from(Atom::from("$gen_call")),
                Term::from(Tuple::from(vec![
                    Term::from(test_pid()),
                    call_ref.to_term(),
                ])),
                Term::from(Atom::from("call")),
            ])),
        )
    }

    fn gen_reply(call_ref: CallRef, result: Term) -> Message {
        Message::send(
            test_pid(),
            Term::from(Tuple::from(vec![call_ref.to_term(), result])),
        )
    }

    /// Spawns a connection task wired to an in-memory peer, returning a
    /// connection to drive it with and the peer's own message channel.
    fn spawn_test_connection() -> (
        RequestSender,
        Sender<ClonableStream<tokio::io::DuplexStream>>,
        Receiver<ClonableStream<tokio::io::DuplexStream>>,
    ) {
        let (local, remote) = tokio::io::duplex(64 * 1024);
        let (sender, receiver) =
            message::channel(ClonableStream::new(local), DistributionFlags::default());
        let (peer_tx, peer_rx) =
            message::channel(ClonableStream::new(remote), DistributionFlags::default());

        let (request_tx, request_rx) = mpsc::channel(32);
        let state = Arc::new(RwLock::new(ConnectionState::Connected));
        let mut conn = NodeConnection::new("peer@localhost".to_string());
        tokio::spawn(connection_task(
            sender,
            receiver,
            request_rx,
            state,
            "peer@localhost".to_string(),
        ));

        conn.request_tx = Some(request_tx);

        (conn.request_sender().unwrap(), peer_tx, peer_rx)
    }

    #[tokio::test]
    async fn reply_to_an_abandoned_call_is_not_delivered_to_a_later_caller() {
        let (conn, mut peer_tx, mut peer_rx) = spawn_test_connection();

        let abandoned_ref = CallRef::next();
        let abandoned = tokio::time::timeout(
            Duration::from_millis(50),
            conn.call(gen_call(abandoned_ref), abandoned_ref),
        )
        .await;
        assert!(abandoned.is_err(), "first call should time out");
        peer_rx.recv().await.unwrap();

        let live_ref = CallRef::next();
        let caller = {
            let conn = conn.clone();
            tokio::spawn(async move { conn.call(gen_call(live_ref), live_ref).await })
        };
        peer_rx.recv().await.unwrap();

        peer_tx
            .send(gen_reply(
                abandoned_ref,
                Term::from(Atom::from("first_result")),
            ))
            .await
            .unwrap();
        peer_tx
            .send(gen_reply(live_ref, Term::from(Atom::from("second_result"))))
            .await
            .unwrap();

        let result = caller.await.unwrap().unwrap();
        assert_eq!(result, Term::from(Atom::from("second_result")));
    }

    #[tokio::test]
    async fn concurrent_calls_each_receive_their_own_reply() {
        let (conn, mut peer_tx, mut peer_rx) = spawn_test_connection();

        let first_ref = CallRef::next();
        let first = {
            let conn = conn.clone();
            tokio::spawn(async move { conn.call(gen_call(first_ref), first_ref).await })
        };
        peer_rx.recv().await.unwrap();

        let second_ref = CallRef::next();
        let second = {
            let conn = conn.clone();
            tokio::spawn(async move { conn.call(gen_call(second_ref), second_ref).await })
        };
        peer_rx.recv().await.unwrap();

        peer_tx
            .send(gen_reply(
                second_ref,
                Term::from(Atom::from("second_result")),
            ))
            .await
            .unwrap();
        peer_tx
            .send(gen_reply(first_ref, Term::from(Atom::from("first_result"))))
            .await
            .unwrap();

        assert_eq!(
            first.await.unwrap().unwrap(),
            Term::from(Atom::from("first_result"))
        );
        assert_eq!(
            second.await.unwrap().unwrap(),
            Term::from(Atom::from("second_result"))
        );
    }

    #[tokio::test]
    async fn call_ref_survives_the_peer_re_encoding_the_id() {
        let (conn, mut peer_tx, mut peer_rx) = spawn_test_connection();

        let call_ref = CallRef::next();
        let caller = {
            let conn = conn.clone();
            tokio::spawn(async move { conn.call(gen_call(call_ref), call_ref).await })
        };
        peer_rx.recv().await.unwrap();

        let small_int_tag = Term::from(Tuple::from(vec![
            Term::from(Atom::from(CALL_REF_TAG)),
            Term::from(FixInteger::from(call_ref.0 as i32)),
        ]));
        peer_tx
            .send(Message::send(
                test_pid(),
                Term::from(Tuple::from(vec![
                    small_int_tag,
                    Term::from(Atom::from("result")),
                ])),
            ))
            .await
            .unwrap();

        assert_eq!(
            caller.await.unwrap().unwrap(),
            Term::from(Atom::from("result"))
        );
    }

    #[tokio::test]
    async fn sweep_drops_calls_whose_caller_is_gone() {
        let mut pending: HashMap<CallRef, mpsc::Sender<RpcResponse>> = HashMap::new();

        let abandoned_ref = CallRef::next();
        let (abandoned_tx, abandoned_rx) = mpsc::channel(1);
        pending.insert(abandoned_ref, abandoned_tx);
        drop(abandoned_rx);

        let live_ref = CallRef::next();
        let (live_tx, _live_rx) = mpsc::channel(1);
        pending.insert(live_ref, live_tx);

        sweep_abandoned(&mut pending);

        assert_eq!(pending.len(), 1);
        assert!(pending.contains_key(&live_ref));
    }

    #[test]
    fn parse_node_name_valid() {
        let (name, host) = parse_node_name("foo@localhost").unwrap();
        assert_eq!(name, "foo");
        assert_eq!(host, "localhost");
    }

    #[test]
    fn parse_node_name_with_domain() {
        let (name, host) = parse_node_name("myapp@server.example.com").unwrap();
        assert_eq!(name, "myapp");
        assert_eq!(host, "server.example.com");
    }

    #[test]
    fn parse_node_name_invalid_no_at() {
        let result = parse_node_name("foobar");
        assert!(result.is_err());
    }

    #[test]
    fn parse_node_name_invalid_multiple_at() {
        let result = parse_node_name("foo@bar@baz");
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn connection_state_display() {
        assert_eq!(ConnectionState::Connecting.to_string(), "connecting");
        assert_eq!(ConnectionState::Connected.to_string(), "connected");
        assert_eq!(ConnectionState::Disconnected.to_string(), "disconnected");
    }

    #[tokio::test]
    async fn node_connection_initial_state() {
        let conn = NodeConnection::new("test@localhost".to_string());
        assert_eq!(conn.node_name, "test@localhost");
        assert_eq!(conn.state().await, ConnectionState::Connecting);
        assert!(conn.connected_at.is_none());
    }
}