netconf_rust/

session.rs

use std::collections::HashMap;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::atomic::{AtomicU8, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use bytes::{Buf, Bytes, BytesMut};
use log::{debug, trace, warn};
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, ReadHalf, WriteHalf};
use tokio_stream::StreamExt;
use tokio_util::codec::{Encoder, FramedRead};

use crate::codec::{DecodedFrame, FramingMode, NetconfCodec, extract_message_id_from_bytes};
use crate::config::Config;
use crate::error::TransportError;
use crate::hello::ServerHello;
use crate::message::{self, DataPayload, RpcReply, RpcReplyBody, ServerMessage};
use crate::stream::NetconfStream;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionState {
    /// Hello exchange complete, ready for RPCs
    Ready = 0,
    /// A 'close-session' RPC has been sent, awaiting reply
    Closing = 1,
    /// Session terminated gracefully or with error
    Closed = 2,
}

impl SessionState {
    fn from_u8(v: u8) -> Self {
        match v {
            0 => Self::Ready,
            1 => Self::Closing,
            _ => Self::Closed,
        }
    }
}

impl std::fmt::Display for SessionState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Ready => write!(f, "Ready"),
            Self::Closing => write!(f, "Closing"),
            Self::Closed => write!(f, "Closed"),
        }
    }
}

/// Reason the session disconnected.
///
/// Delivered via [`Session::disconnected()`] when the background reader
/// task detects that the connection is no longer alive.
#[derive(Debug, Clone)]
pub enum DisconnectReason {
    /// The remote end closed the connection cleanly (TCP FIN / EOF).
    Eof,
    /// A transport error severed the connection.
    ///
    /// Contains the error's display string.
    TransportError(String),
    /// The [`Session`] was dropped without calling [`Session::close_session()`].
    Dropped,
}

impl std::fmt::Display for DisconnectReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Eof => write!(f, "connection closed by remote"),
            Self::TransportError(e) => write!(f, "transport error: {e}"),
            Self::Dropped => write!(f, "session dropped"),
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub enum Datastore {
    Running,
    Candidate,
    Startup,
}

impl Datastore {
    fn as_xml(&self) -> &'static str {
        match self {
            Datastore::Running => "<running/>",
            Datastore::Candidate => "<candidate/>",
            Datastore::Startup => "<startup/>",
        }
    }
}

// =============================================================================
// PendingRpc — supports both normal (buffered) and streaming RPCs
// =============================================================================

/// A pending RPC awaiting its response from the server.
///
/// Normal RPCs accumulate the full response and deliver it as an `RpcReply`
/// via a oneshot channel. Streaming RPCs forward individual chunks via an
/// mpsc channel, enabling the consumer to process data incrementally.
enum PendingRpc {
    Normal(tokio::sync::oneshot::Sender<crate::Result<RpcReply>>),
    Stream(tokio::sync::mpsc::Sender<crate::Result<Bytes>>),
}

/// A handle to a pending RPC reply.
///
/// Created by [`Session::rpc_send()`], this allows sending multiple RPCs
/// before awaiting any replies (pipelining).
pub struct RpcFuture {
    rx: tokio::sync::oneshot::Receiver<crate::Result<RpcReply>>,
    msg_id: u32,
    rpc_timeout: Option<Duration>,
}

impl RpcFuture {
    /// The message-id of this RPC.
    pub fn message_id(&self) -> u32 {
        self.msg_id
    }

    /// Await the RPC reply.
    ///
    /// If an `rpc_timeout` was configured, this will fail with
    /// [`TransportError::Timeout`]
    /// if the server does not reply in time.
    pub async fn response(self) -> crate::Result<RpcReply> {
        let result = match self.rpc_timeout {
            Some(duration) => tokio::time::timeout(duration, self.rx)
                .await
                .map_err(|_| crate::Error::Transport(TransportError::Timeout(duration)))?,
            None => self.rx.await,
        };
        result.map_err(|_| crate::Error::SessionClosed)?
    }

    /// Await the RPC reply with an explicit timeout, ignoring the
    /// session-level `rpc_timeout`.
    ///
    /// Fails with [`TransportError::Timeout`]
    /// if the server does not reply within `timeout`.
    pub async fn response_with_timeout(self, timeout: Duration) -> crate::Result<RpcReply> {
        let result = tokio::time::timeout(timeout, self.rx)
            .await
            .map_err(|_| crate::Error::Transport(TransportError::Timeout(timeout)))?;
        result.map_err(|_| crate::Error::SessionClosed)?
    }
}

// =============================================================================
// RpcStream — streaming RPC response with AsyncRead
// =============================================================================

/// A streaming RPC response that yields raw XML bytes as they arrive.
///
/// Implements [`AsyncRead`] so it can be plugged directly into compression
/// encoders, file writers, or any byte-oriented consumer without buffering
/// the entire response.
///
/// Created by [`Session::rpc_stream()`].
///
/// # Example
///
/// ```no_run
/// # use netconf_rust::Session;
/// # use tokio::io::AsyncReadExt;
/// # async fn example(session: &Session) -> netconf_rust::Result<()> {
/// let mut stream = session.rpc_stream("<get-config><source><running/></source></get-config>").await?;
/// let mut buf = [0u8; 8192];
/// loop {
///     let n = stream.read(&mut buf).await?;
///     if n == 0 { break; }
///     // process buf[..n] — e.g. feed to a compressor
/// }
/// # Ok(())
/// # }
/// ```
pub struct RpcStream {
    rx: tokio::sync::mpsc::Receiver<crate::Result<Bytes>>,
    /// Partially consumed chunk from the last `poll_read`.
    current: Bytes,
    msg_id: u32,
    done: bool,
}

impl RpcStream {
    /// The message-id of this streaming RPC.
    pub fn message_id(&self) -> u32 {
        self.msg_id
    }

    /// Whether the stream has finished (all chunks received).
    pub fn is_done(&self) -> bool {
        self.done
    }
}

impl AsyncRead for RpcStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // 1. Serve from the partially-consumed current chunk.
        if !self.current.is_empty() {
            let n = std::cmp::min(buf.remaining(), self.current.len());
            buf.put_slice(&self.current[..n]);
            self.current.advance(n);
            return Poll::Ready(Ok(()));
        }

        // 2. If done, return EOF.
        if self.done {
            return Poll::Ready(Ok(()));
        }

        // 3. Poll the channel for the next chunk.
        match self.rx.poll_recv(cx) {
            Poll::Ready(Some(Ok(chunk))) => {
                let n = std::cmp::min(buf.remaining(), chunk.len());
                buf.put_slice(&chunk[..n]);
                if n < chunk.len() {
                    self.current = chunk.slice(n..);
                }
                Poll::Ready(Ok(()))
            }
            Poll::Ready(Some(Err(e))) => {
                self.done = true;
                Poll::Ready(Err(io::Error::other(e.to_string())))
            }
            Poll::Ready(None) => {
                // Channel closed — stream complete.
                self.done = true;
                Poll::Ready(Ok(()))
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

// =============================================================================
// SessionInner
// =============================================================================

struct SessionInner {
    /// Pending RPC replies: message-id → pending RPC (normal or stream).
    pending: Mutex<HashMap<u32, PendingRpc>>,
    /// Session lifecycle state, stored as AtomicU8 for lock-free access
    /// across the session and reader task.
    ///
    /// State transitions:
    ///   Ready → Closing → Closed   (graceful: user calls close_session)
    ///   Ready → Closed             (abrupt: reader hits EOF or error)
    ///
    /// - Ready:   normal operation, RPCs can be sent
    /// - Closing: close-session RPC in flight, new RPCs rejected
    /// - Closed:  session terminated, all operations fail
    state: AtomicU8,
    /// Sender for the disconnect notification. The reader_loop sends the
    /// reason just before exiting. Receivers (from `Session::disconnected()`)
    /// wake up immediately.
    disconnect_tx: tokio::sync::watch::Sender<Option<DisconnectReason>>,
    /// Anchor instant for computing `last_rpc_at` from `last_rpc_nanos`.
    created_at: Instant,
    /// Nanoseconds elapsed since `created_at` when the last RPC reply was
    /// successfully routed. 0 means no reply has been received yet.
    last_rpc_nanos: AtomicU64,
    /// Number of streaming RPCs currently in flight (removed from `pending`
    /// but not yet completed with `EndOfMessage`).
    active_streams: AtomicUsize,
}

impl SessionInner {
    fn state(&self) -> SessionState {
        SessionState::from_u8(self.state.load(Ordering::Acquire))
    }

    fn set_state(&self, state: SessionState) {
        self.state.store(state as u8, Ordering::Release);
    }

    fn drain_pending(&self) -> usize {
        let mut pending = self.pending.lock().unwrap();
        let count = pending.len();
        for (_, rpc) in pending.drain() {
            match rpc {
                PendingRpc::Normal(tx) => {
                    let _ = tx.send(Err(crate::Error::SessionClosed));
                }
                PendingRpc::Stream(tx) => {
                    let _ = tx.try_send(Err(crate::Error::SessionClosed));
                }
            }
        }
        count
    }
}

/// Write half + codec, behind a `tokio::sync::Mutex` so Session methods
/// can take `&self`. The Mutex is held only during encode + write + flush.
struct WriterState {
    writer: WriteHalf<NetconfStream>,
    codec: NetconfCodec,
}

/// A NETCONF session with pipelining and streaming support.
///
/// Architecture:
/// - **Writer**: the session holds the write half of the stream and a codec
///   for encoding outgoing RPCs.
/// - **Reader task**: a background tokio task reads framed chunks from
///   the read half, classifies them (`<rpc-reply>` vs `<notification>`),
///   and routes them to the correct handler.
/// - **Pipelining**: [`rpc_send()`](Session::rpc_send) writes an RPC and
///   returns an [`RpcFuture`] without waiting for the reply. Multiple RPCs
///   can be in flight simultaneously.
/// - **Streaming**: [`rpc_stream()`](Session::rpc_stream) writes an RPC and
///   returns an [`RpcStream`] that yields raw bytes as they arrive. This
///   avoids buffering the entire response for large payloads.
pub struct Session {
    /// Write state behind an async Mutex. Held only for the duration of
    /// encoding and flushing a single message.
    writer: tokio::sync::Mutex<WriterState>,

    /// Shared state between this session and the background reader task.
    /// Contains the pending RPC map, notification channel, and session state.
    inner: Arc<SessionInner>,

    /// The server's hello response, containing its capabilities and session ID.
    server_hello: ServerHello,

    /// Negotiated framing mode (EOM for 1.0-only servers, chunked for 1.1).
    framing: FramingMode,

    /// Timeout applied to each RPC response wait.
    rpc_timeout: Option<Duration>,

    /// Channel capacity for streaming RPCs (see [`Config::stream_buffer_capacity`]).
    stream_buffer_capacity: usize,

    /// Receiver for disconnect notifications. Cloned for each call to
    /// `disconnected()`, allowing multiple independent subscribers.
    disconnect_rx: tokio::sync::watch::Receiver<Option<DisconnectReason>>,

    /// Instant when the session was established (hello exchange complete).
    connected_since: Instant,

    /// Handle to the background reader task. Aborted on drop to ensure
    /// the task doesn't outlive the session. Declared before `_keep_alive`
    /// so it is dropped first — the reader task must be aborted before
    /// the SSH connection is torn down.
    _reader_handle: tokio::task::JoinHandle<()>,

    /// Holds the SSH handle alive. Dropping this tears down the SSH connection,
    /// which would invalidate the stream. Never accessed, just kept alive.
    _keep_alive: Option<Box<dyn std::any::Any + Send + Sync>>,
}

impl Drop for Session {
    fn drop(&mut self) {
        // 1. Drain pending RPCs → waiters get deterministic SessionClosed
        let drained = self.inner.drain_pending();
        if drained > 0 {
            debug!(
                "session {}: drop: drained {drained} pending RPCs",
                self.server_hello.session_id
            );
        }
        // 2. Mark session closed
        self.inner.set_state(SessionState::Closed);
        // 3. Notify disconnect subscribers (only if reader_loop hasn't already)
        self.inner.disconnect_tx.send_if_modified(|current| {
            if current.is_none() {
                *current = Some(DisconnectReason::Dropped);
                true
            } else {
                false
            }
        });
        // 4. Abort reader task (its cleanup is now a no-op)
        self._reader_handle.abort();
    }
}

impl Session {
    /// Connect to a NETCONF server over SSH with password authentication.
    pub async fn connect(
        host: &str,
        port: u16,
        username: &str,
        password: &str,
    ) -> crate::Result<Self> {
        Self::connect_with_config(host, port, username, password, Config::default()).await
    }

    /// Connect with custom configuration.
    pub async fn connect_with_config(
        host: &str,
        port: u16,
        username: &str,
        password: &str,
        config: Config,
    ) -> crate::Result<Self> {
        let (mut stream, keep_alive) =
            crate::transport::connect(host, port, username, password, &config).await?;
        let (server_hello, framing) = exchange_hello(&mut stream, &config).await?;
        Self::build(stream, Some(keep_alive), server_hello, framing, config)
    }

    /// Create a session from an existing stream (useful for testing).
    pub async fn from_stream<S: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
        stream: S,
    ) -> crate::Result<Self> {
        Self::from_stream_with_config(stream, Config::default()).await
    }

    /// Create a session from an existing stream with custom configuration.
    pub async fn from_stream_with_config<S: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
        mut stream: S,
        config: Config,
    ) -> crate::Result<Self> {
        let (server_hello, framing) = exchange_hello(&mut stream, &config).await?;
        let boxed: NetconfStream = Box::new(stream);
        Self::build(boxed, None, server_hello, framing, config)
    }

    fn build(
        stream: NetconfStream,
        keep_alive: Option<Box<dyn std::any::Any + Send + Sync>>,
        server_hello: ServerHello,
        framing: FramingMode,
        config: Config,
    ) -> crate::Result<Self> {
        debug!(
            "session {}: building (framing={:?}, capabilities={})",
            server_hello.session_id,
            framing,
            server_hello.capabilities.len()
        );
        let (read_half, write_half) = tokio::io::split(stream);

        let read_codec = NetconfCodec::new(framing, config.codec);
        let write_codec = NetconfCodec::new(framing, config.codec);
        let reader = FramedRead::new(read_half, read_codec);

        let (disconnect_tx, disconnect_rx) = tokio::sync::watch::channel(None);

        let inner = Arc::new(SessionInner {
            pending: Mutex::new(HashMap::new()),
            state: AtomicU8::new(SessionState::Ready as u8),
            disconnect_tx,
            created_at: Instant::now(),
            last_rpc_nanos: AtomicU64::new(0),
            active_streams: AtomicUsize::new(0),
        });

        let reader_inner = Arc::clone(&inner);
        let session_id = server_hello.session_id;
        let reader_handle = tokio::spawn(async move {
            reader_loop(reader, reader_inner, session_id).await;
        });

        Ok(Self {
            writer: tokio::sync::Mutex::new(WriterState {
                writer: write_half,
                codec: write_codec,
            }),
            inner,
            server_hello,
            framing,
            rpc_timeout: config.rpc_timeout,
            stream_buffer_capacity: config.stream_buffer_capacity,
            disconnect_rx,
            connected_since: Instant::now(),
            _reader_handle: reader_handle,
            _keep_alive: keep_alive,
        })
    }

    pub fn session_id(&self) -> u32 {
        self.server_hello.session_id
    }

    pub fn server_capabilities(&self) -> &[String] {
        &self.server_hello.capabilities
    }

    pub fn framing_mode(&self) -> FramingMode {
        self.framing
    }

    pub fn state(&self) -> SessionState {
        self.inner.state()
    }

    /// Returns a future that completes when the session disconnects.
    ///
    /// Can be called multiple times — each call clones an internal
    /// `watch::Receiver`, so multiple tasks can independently await
    /// the same disconnect event. If the session is already disconnected
    /// when called, returns immediately.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use netconf_rust::Session;
    /// # async fn example(session: &Session) {
    /// let reason = session.disconnected().await;
    /// println!("session died: {reason}");
    /// # }
    /// ```
    pub fn disconnected(&self) -> impl Future<Output = DisconnectReason> + Send + 'static {
        let mut rx = self.disconnect_rx.clone();
        async move {
            // Check if already disconnected (late subscriber).
            if let Some(reason) = rx.borrow_and_update().clone() {
                return reason;
            }
            // Wait for the reader_loop to send the reason. If the sender
            // is dropped (Session dropped, reader aborted), treat as Dropped.
            loop {
                if rx.changed().await.is_err() {
                    return DisconnectReason::Dropped;
                }
                if let Some(reason) = rx.borrow_and_update().clone() {
                    return reason;
                }
            }
        }
    }

    fn check_state(&self) -> crate::Result<()> {
        let state = self.inner.state();
        if state != SessionState::Ready {
            return Err(crate::Error::InvalidState(state.to_string()));
        }
        Ok(())
    }

    /// Encode a message with the negotiated framing (EOM or chunked) and
    /// write it to the stream.
    async fn send_encoded(&self, xml: &str) -> crate::Result<()> {
        let mut buf = BytesMut::new();
        let mut state = self.writer.lock().await;
        state.codec.encode(Bytes::from(xml.to_string()), &mut buf)?;
        trace!(
            "session {}: writing {} bytes to stream",
            self.server_hello.session_id,
            buf.len()
        );
        state.writer.write_all(&buf).await?;
        state.writer.flush().await?;
        Ok(())
    }

    /// Send a raw RPC and return a future for the reply (pipelining).
    ///
    /// This writes the RPC to the server immediately but does not wait
    /// for the reply. Call [`RpcFuture::response()`] to await the reply.
    /// Multiple RPCs can be pipelined by calling this repeatedly before
    /// awaiting any of them.
    pub async fn rpc_send(&self, inner_xml: &str) -> crate::Result<RpcFuture> {
        self.check_state()?;
        let (msg_id, xml) = message::build_rpc(inner_xml);
        debug!(
            "session {}: sending rpc message-id={} ({} bytes)",
            self.server_hello.session_id,
            msg_id,
            xml.len()
        );
        trace!(
            "session {}: rpc content: {}",
            self.server_hello.session_id, inner_xml
        );
        let (tx, rx) = tokio::sync::oneshot::channel();

        self.inner
            .pending
            .lock()
            .unwrap()
            .insert(msg_id, PendingRpc::Normal(tx));

        if let Err(e) = self.send_encoded(&xml).await {
            debug!(
                "session {}: send failed for message-id={}: {}",
                self.server_hello.session_id, msg_id, e
            );
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }
        Ok(RpcFuture {
            rx,
            msg_id,
            rpc_timeout: self.rpc_timeout,
        })
    }

    /// Send a raw RPC and wait for the reply.
    pub async fn rpc_raw(&self, inner_xml: &str) -> crate::Result<RpcReply> {
        let future = self.rpc_send(inner_xml).await?;
        future.response().await
    }

    /// Send a raw RPC and return a streaming response.
    ///
    /// Unlike [`rpc_send()`](Session::rpc_send) which buffers the entire
    /// response, this returns an [`RpcStream`] that yields raw XML bytes
    /// as individual chunks arrive from the server. The stream implements
    /// [`AsyncRead`], so it can be piped directly into compression encoders,
    /// file writers, or any byte-oriented consumer.
    ///
    /// Normal (buffered) and streaming RPCs can be freely interleaved on
    /// the same session — the reader task routes each message independently
    /// based on its `message-id`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use netconf_rust::Session;
    /// # use tokio::io::AsyncReadExt;
    /// # async fn example(session: &Session) -> netconf_rust::Result<()> {
    /// // Pipeline a normal RPC alongside a streaming one
    /// let edit_future = session.rpc_send("<edit-config>...</edit-config>").await?;
    /// let mut config_stream = session.rpc_stream("<get-config><source><running/></source></get-config>").await?;
    ///
    /// // Stream the large response
    /// let mut buf = [0u8; 8192];
    /// loop {
    ///     let n = config_stream.read(&mut buf).await?;
    ///     if n == 0 { break; }
    ///     // process buf[..n]
    /// }
    ///
    /// // Collect the normal RPC reply
    /// let edit_reply = edit_future.response().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn rpc_stream(&self, inner_xml: &str) -> crate::Result<RpcStream> {
        self.check_state()?;
        let (msg_id, xml) = message::build_rpc(inner_xml);
        debug!(
            "session {}: sending streaming rpc message-id={} ({} bytes)",
            self.server_hello.session_id,
            msg_id,
            xml.len()
        );

        let (tx, rx) = tokio::sync::mpsc::channel(self.stream_buffer_capacity);

        self.inner
            .pending
            .lock()
            .unwrap()
            .insert(msg_id, PendingRpc::Stream(tx));

        if let Err(e) = self.send_encoded(&xml).await {
            debug!(
                "session {}: send failed for streaming message-id={}: {}",
                self.server_hello.session_id, msg_id, e
            );
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }

        Ok(RpcStream {
            rx,
            current: Bytes::new(),
            msg_id,
            done: false,
        })
    }

    /// Internal rpc send that skips state check. Only used for sending close-session.
    async fn rpc_send_unchecked(&self, inner_xml: &str) -> crate::Result<RpcFuture> {
        let (msg_id, xml) = message::build_rpc(inner_xml);
        let (tx, rx) = tokio::sync::oneshot::channel();

        self.inner
            .pending
            .lock()
            .unwrap()
            .insert(msg_id, PendingRpc::Normal(tx));

        if let Err(e) = self.send_encoded(&xml).await {
            self.inner.pending.lock().unwrap().remove(&msg_id);
            return Err(e);
        }

        Ok(RpcFuture {
            rx,
            msg_id,
            rpc_timeout: self.rpc_timeout,
        })
    }

    /// Retrieve configuration from a datastore.
    pub async fn get_config(
        &self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve configuration as a zero-copy `DataPayload`.
    ///
    /// Same as `get_config()` but returns a `DataPayload` instead of `String`,
    /// avoiding a copy of the response body. Use `payload.as_str()` for a
    /// zero-copy `&str` view, or `payload.reader()` for streaming XML events.
    pub async fn get_config_payload(
        &self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Retrieve configuration from a datastore as a streaming response.
    ///
    /// Same as `get_config()` but returns an [`RpcStream`] instead of
    /// buffering the entire response.
    pub async fn get_config_stream(
        &self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<RpcStream> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        self.rpc_stream(&inner).await
    }

    /// Retrieve running configuration and state data.
    pub async fn get(&self, filter: Option<&str>) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve running configuration and state data as a zero-copy `DataPayload`.
    ///
    /// Same as `get()` but returns a `DataPayload` instead of `String`,
    /// avoiding a copy of the response body.
    pub async fn get_payload(&self, filter: Option<&str>) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Retrieve running configuration and state data as a streaming response.
    ///
    /// Same as `get()` but returns an [`RpcStream`] instead of buffering
    /// the entire response.
    pub async fn get_stream(&self, filter: Option<&str>) -> crate::Result<RpcStream> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        self.rpc_stream(&inner).await
    }

    /// Edit the configuration of a target datastore.
    pub async fn edit_config(&self, target: Datastore, config: &str) -> crate::Result<()> {
        let inner = format!(
            "<edit-config><target>{}</target><config>{config}</config></edit-config>",
            target.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Lock a datastore
    pub async fn lock(&self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<lock><target>{}</target></lock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Unlock a datastore.
    pub async fn unlock(&self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<unlock><target>{}</target></unlock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Commit the candidate configuration to running.
    pub async fn commit(&self) -> crate::Result<()> {
        let reply = self.rpc_raw("<commit/>").await?;
        reply_to_ok(reply)
    }

    /// Gracefully close the NETCONF session.
    pub async fn close_session(&self) -> crate::Result<()> {
        // Atomically transition Ready → Closing. If another caller already
        // moved us out of Ready, we fail immediately.
        let prev = self.inner.state.compare_exchange(
            SessionState::Ready as u8,
            SessionState::Closing as u8,
            Ordering::AcqRel,
            Ordering::Acquire,
        );
        if let Err(current) = prev {
            let state = SessionState::from_u8(current);
            return Err(crate::Error::InvalidState(state.to_string()));
        }
        debug!("session {}: closing", self.server_hello.session_id);
        let result = self.rpc_send_unchecked("<close-session/>").await;
        match result {
            Ok(future) => {
                let reply = future.response().await;
                self.inner.set_state(SessionState::Closed);
                debug!(
                    "session {}: closed gracefully",
                    self.server_hello.session_id
                );
                reply_to_ok(reply?)
            }
            Err(e) => {
                self.inner.set_state(SessionState::Closed);
                debug!(
                    "session {}: close failed: {}",
                    self.server_hello.session_id, e
                );
                Err(e)
            }
        }
    }

    /// Gracefully close the NETCONF session and shut down the transport.
    ///
    /// Sends a `<close-session/>` RPC, shuts down the write half of the
    /// stream, and drops the session (aborting the reader task and releasing
    /// the SSH handle). Prefer this over [`close_session()`](Self::close_session)
    /// when you are done with the session entirely.
    pub async fn close(self) -> crate::Result<()> {
        let result = self.close_session().await;
        self.writer.lock().await.writer.shutdown().await.ok();
        // `self` is dropped here — reader task aborted, SSH handle released
        result
    }

    /// Force-close another NETCONF session.
    pub async fn kill_session(&self, session_id: u32) -> crate::Result<()> {
        let inner = format!("<kill-session><session-id>{session_id}</session-id></kill-session>");
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Return a wrapper that applies `timeout` to every RPC sent through it,
    /// overriding the session-level `rpc_timeout`.
    pub fn with_timeout(&self, timeout: Duration) -> SessionWithTimeout<'_> {
        SessionWithTimeout {
            session: self,
            timeout,
        }
    }

    /// Number of RPCs that have been sent but not yet fully replied to.
    ///
    /// This includes both RPCs awaiting their first reply byte (in the
    /// pending map) and streaming RPCs whose response is in flight but
    /// not yet complete.
    pub fn pending_rpc_count(&self) -> usize {
        self.inner.pending.lock().unwrap().len() + self.inner.active_streams.load(Ordering::Acquire)
    }

    /// The [`Instant`] when the most recent RPC reply was received, or
    /// `None` if no reply has been received yet.
    pub fn last_rpc_at(&self) -> Option<Instant> {
        let nanos = self.inner.last_rpc_nanos.load(Ordering::Acquire);
        if nanos == 0 {
            None
        } else {
            Some(self.inner.created_at + Duration::from_nanos(nanos))
        }
    }

    /// The [`Instant`] when the session was established (hello exchange
    /// complete and reader task started).
    pub fn connected_since(&self) -> Instant {
        self.connected_since
    }
}

/// A wrapper around [`Session`] that applies a per-call timeout to every RPC.
///
/// Created by [`Session::with_timeout()`]. Each method sends the RPC via
/// the underlying session and awaits the reply with the configured timeout.
pub struct SessionWithTimeout<'a> {
    session: &'a Session,
    timeout: Duration,
}

impl SessionWithTimeout<'_> {
    /// Send a raw RPC and wait for the reply with the configured timeout.
    pub async fn rpc_raw(&self, inner_xml: &str) -> crate::Result<RpcReply> {
        let future = self.session.rpc_send(inner_xml).await?;
        future.response_with_timeout(self.timeout).await
    }

    /// Retrieve configuration from a datastore.
    pub async fn get_config(
        &self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve configuration as a zero-copy `DataPayload`.
    pub async fn get_config_payload(
        &self,
        source: Datastore,
        filter: Option<&str>,
    ) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!(
            "<get-config><source>{}</source>{filter_xml}</get-config>",
            source.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Retrieve running configuration and state data.
    pub async fn get(&self, filter: Option<&str>) -> crate::Result<String> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply_to_data(reply)
    }

    /// Retrieve running configuration and state data as a zero-copy `DataPayload`.
    pub async fn get_payload(&self, filter: Option<&str>) -> crate::Result<DataPayload> {
        let filter_xml = match filter {
            Some(f) => format!(r#"<filter type="subtree">{f}</filter>"#),
            None => String::new(),
        };
        let inner = format!("<get>{filter_xml}</get>");
        let reply = self.rpc_raw(&inner).await?;
        reply.into_data()
    }

    /// Edit the configuration of a target datastore.
    pub async fn edit_config(&self, target: Datastore, config: &str) -> crate::Result<()> {
        let inner = format!(
            "<edit-config><target>{}</target><config>{config}</config></edit-config>",
            target.as_xml()
        );
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Lock a datastore.
    pub async fn lock(&self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<lock><target>{}</target></lock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Unlock a datastore.
    pub async fn unlock(&self, target: Datastore) -> crate::Result<()> {
        let inner = format!("<unlock><target>{}</target></unlock>", target.as_xml());
        let reply = self.rpc_raw(&inner).await?;
        reply_to_ok(reply)
    }

    /// Commit the candidate configuration to running.
    pub async fn commit(&self) -> crate::Result<()> {
        let reply = self.rpc_raw("<commit/>").await?;
        reply_to_ok(reply)
    }
}

/// Perform the NETCONF hello exchange, optionally with a timeout.
async fn exchange_hello<S: AsyncRead + AsyncWrite + Unpin>(
    stream: &mut S,
    config: &Config,
) -> crate::Result<(ServerHello, FramingMode)> {
    let fut = crate::hello::exchange(stream, config.codec.max_message_size);
    match config.hello_timeout {
        Some(duration) => tokio::time::timeout(duration, fut)
            .await
            .map_err(|_| crate::Error::Transport(TransportError::Timeout(duration)))?,
        None => fut.await,
    }
}

// =============================================================================
// Reader task — state machine for chunk-level routing
// =============================================================================

/// Per-message state machine for the reader task.
///
/// Tracks whether the current message is being accumulated (normal RPC)
/// or streamed (streaming RPC). The state transitions on each
/// [`DecodedFrame`] from the codec.
enum ReaderMessageState {
    /// Waiting for the first chunk of a new message.
    /// Accumulates bytes until the `message-id` can be extracted.
    AwaitingHeader { buf: BytesMut },
    /// Accumulating chunks for a normal (non-streaming) RPC.
    Accumulating { msg_id: u32, buf: BytesMut },
    /// Forwarding chunks for a streaming RPC.
    Streaming {
        msg_id: u32,
        tx: tokio::sync::mpsc::Sender<crate::Result<Bytes>>,
    },
}

/// This loop is the only thing reading from the SSH stream. It runs in a
/// background tokio task. The session's main API (the writer side) never
/// reads — it only writes RPCs and waits on oneshot/mpsc channels.
/// This separation is what makes pipelining and streaming work.
///
/// The codec now yields [`DecodedFrame`] items (individual chunks and
/// end-of-message markers) instead of complete messages. The reader task
/// uses a [`ReaderMessageState`] state machine to decide per-message
/// whether to accumulate chunks (normal RPCs) or forward them (streaming
/// RPCs). The `message-id` from the first chunk determines the routing.
async fn reader_loop(
    mut reader: FramedRead<ReadHalf<NetconfStream>, NetconfCodec>,
    inner: Arc<SessionInner>,
    session_id: u32,
) {
    debug!("session {}: reader loop started", session_id);
    let mut disconnect_reason = DisconnectReason::Eof;
    let mut state = ReaderMessageState::AwaitingHeader {
        buf: BytesMut::new(),
    };

    loop {
        // Propagate the closing flag so decode_eof can discard holdback bytes.
        if inner.state() == SessionState::Closing {
            reader.decoder_mut().set_closing();
        }
        let Some(result) = reader.next().await else {
            break;
        };
        match result {
            Ok(frame) => {
                state = process_frame(frame, state, &inner, session_id).await;
            }
            Err(e) => {
                debug!("session {}: reader error: {e}", session_id);
                disconnect_reason = DisconnectReason::TransportError(e.to_string());

                // Notify any in-flight streaming RPC about the error.
                if let ReaderMessageState::Streaming { tx, .. } = &state {
                    let _ = tx.try_send(Err(crate::Error::SessionClosed));
                }

                let drained = inner.drain_pending();
                if drained > 0 {
                    debug!(
                        "session {}: drained {} pending RPCs after error",
                        session_id, drained
                    );
                }
                break;
            }
        }
    }

    // Notify any in-flight streaming RPC about session close.
    if let ReaderMessageState::Streaming { tx, .. } = &state {
        let _ = tx.try_send(Err(crate::Error::SessionClosed));
    }

    // Drain any remaining pending RPCs.
    {
        let drained = inner.drain_pending();
        if drained > 0 {
            debug!(
                "session {}: drained {} pending RPCs on stream close",
                session_id, drained
            );
        }
    }

    inner.set_state(SessionState::Closed);
    let _ = inner.disconnect_tx.send(Some(disconnect_reason));
    debug!("session {}: reader loop ended", session_id);
}

/// Process a single decoded frame, advancing the per-message state machine.
async fn process_frame(
    frame: DecodedFrame,
    state: ReaderMessageState,
    inner: &SessionInner,
    session_id: u32,
) -> ReaderMessageState {
    match frame {
        DecodedFrame::Chunk(chunk) => match state {
            ReaderMessageState::AwaitingHeader { mut buf } => {
                buf.extend_from_slice(&chunk);

                // Try to extract the message-id from what we have so far.
                if let Some(msg_id) = extract_message_id_from_bytes(&buf) {
                    // Look up the pending RPC type to decide routing.
                    let is_stream = {
                        let pending = inner.pending.lock().unwrap();
                        matches!(pending.get(&msg_id), Some(PendingRpc::Stream(_)))
                    };

                    if is_stream {
                        // Remove the stream sender from pending and transition
                        // to Streaming state.
                        let tx = {
                            let mut pending = inner.pending.lock().unwrap();
                            match pending.remove(&msg_id) {
                                Some(PendingRpc::Stream(tx)) => tx,
                                // Race: might have been drained between the
                                // check above and here. Fall back to accumulate.
                                _ => {
                                    return ReaderMessageState::Accumulating { msg_id, buf };
                                }
                            }
                        };
                        inner.active_streams.fetch_add(1, Ordering::Release);
                        // Forward the buffered header as the first chunk.
                        let _ = tx.send(Ok(buf.freeze())).await;
                        debug!(
                            "session {}: streaming rpc message-id={}",
                            session_id, msg_id
                        );
                        ReaderMessageState::Streaming { msg_id, tx }
                    } else {
                        // Normal RPC or unknown — accumulate.
                        ReaderMessageState::Accumulating { msg_id, buf }
                    }
                } else {
                    // Need more data to find message-id.
                    ReaderMessageState::AwaitingHeader { buf }
                }
            }
            ReaderMessageState::Accumulating { msg_id, mut buf } => {
                buf.extend_from_slice(&chunk);
                ReaderMessageState::Accumulating { msg_id, buf }
            }
            ReaderMessageState::Streaming { msg_id, tx } => {
                // Forward chunk to the streaming consumer. If the consumer
                // dropped the receiver, just discard — we still need to
                // consume until EndOfMessage so we can start the next message.
                let _ = tx.send(Ok(chunk)).await;
                ReaderMessageState::Streaming { msg_id, tx }
            }
        },

        DecodedFrame::EndOfMessage => match state {
            ReaderMessageState::AwaitingHeader { .. } => {
                // Empty message or we never found a message-id — reset.
                trace!("session {}: empty or unparseable message", session_id);
                ReaderMessageState::AwaitingHeader {
                    buf: BytesMut::new(),
                }
            }
            ReaderMessageState::Accumulating { msg_id, buf } => {
                // Complete normal message — classify and route.
                let bytes = buf.freeze();
                trace!(
                    "session {}: complete message for msg-id={} ({} bytes)",
                    session_id,
                    msg_id,
                    bytes.len()
                );

                match message::classify_message(bytes) {
                    Ok(ServerMessage::RpcReply(reply)) => {
                        debug!(
                            "session {}: received rpc-reply message-id={}",
                            session_id, reply.message_id
                        );
                        let tx = {
                            let mut pending = inner.pending.lock().unwrap();
                            pending.remove(&reply.message_id)
                        };
                        if let Some(PendingRpc::Normal(tx)) = tx {
                            let nanos = inner.created_at.elapsed().as_nanos() as u64;
                            inner.last_rpc_nanos.store(nanos, Ordering::Release);
                            let _ = tx.send(Ok(reply));
                        } else {
                            warn!(
                                "session {}: received reply for unknown message-id {}",
                                session_id, reply.message_id
                            );
                        }
                    }
                    Err(e) => {
                        warn!("session {}: failed to classify message: {e}", session_id);
                    }
                }

                ReaderMessageState::AwaitingHeader {
                    buf: BytesMut::new(),
                }
            }
            ReaderMessageState::Streaming { msg_id, tx } => {
                // End of streaming message — drop the sender to signal EOF.
                drop(tx);
                inner.active_streams.fetch_sub(1, Ordering::Release);
                let nanos = inner.created_at.elapsed().as_nanos() as u64;
                inner.last_rpc_nanos.store(nanos, Ordering::Release);
                debug!(
                    "session {}: streaming message complete for msg-id={}",
                    session_id, msg_id
                );
                ReaderMessageState::AwaitingHeader {
                    buf: BytesMut::new(),
                }
            }
        },
    }
}

fn reply_to_data(reply: RpcReply) -> crate::Result<String> {
    match reply.body {
        RpcReplyBody::Data(payload) => Ok(payload.into_string()),
        RpcReplyBody::Ok => Ok(String::new()),
        RpcReplyBody::Error(errors) => Err(crate::Error::Rpc {
            message_id: reply.message_id,
            error: errors
                .first()
                .map(|e| e.error_message.clone())
                .unwrap_or_default(),
        }),
    }
}

fn reply_to_ok(reply: RpcReply) -> crate::Result<()> {
    match reply.body {
        RpcReplyBody::Ok => Ok(()),
        RpcReplyBody::Data(_) => Ok(()),
        RpcReplyBody::Error(errors) => Err(crate::Error::Rpc {
            message_id: reply.message_id,
            error: errors
                .first()
                .map(|e| e.error_message.clone())
                .unwrap_or_default(),
        }),
    }
}