// crates/qrusty_client/src/ws.rs
//
// Rust WebSocket client for the Qrusty priority queue server.
//
// Requirements: WS-0017, WS-0020, WS-0022, WS-0023, WS-0024, WS-0025, WS-0026
//
// # Example
//
// ```rust,no_run
// use qrusty_client::ws::{WsSession, DeliveredMessage};
//
// #[tokio::main]
// async fn main() -> anyhow::Result<()> {
//     let session = WsSession::connect("ws://localhost:6784").await?;
//
//     // Publish
//     let id = session.publish("orders", "payload", Some(10u64.into())).await?;
//     println!("published {}", id);
//
//     // Subscribe and process messages
//     let mut rx = session.subscribe("orders").await?;
//     while let Some(Ok(msg)) = rx.recv().await {
//         println!("got {} from {}", msg.payload, msg.queue);
//         session.ack(&msg.queue, &msg.id).await?;
//     }
//
//     session.close().await?;
//     Ok(())
// }
// ```

use crate::error::QrustyClientError;
use crate::priority::Priority;
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, oneshot, Mutex};
use tokio::task::JoinHandle;
use tokio_tungstenite::{connect_async, tungstenite::Message as TMsg};

// ---------------------------------------------------------------------------
// Pending request registry and subscriber registry
// ---------------------------------------------------------------------------

/// Pending request registry: req_id → oneshot sender for the response frame.
type PendingMap = Arc<Mutex<HashMap<String, oneshot::Sender<Result<Value, QrustyClientError>>>>>;

/// Subscriber registry: queue_name → mpsc sender for delivered messages.
type SubMap =
    Arc<Mutex<HashMap<String, mpsc::Sender<Result<DeliveredMessage, QrustyClientError>>>>>;

/// Internal delivery envelope forwarded from router_task to delivery_task.
type DeliveryEnvelope = (String, Result<DeliveredMessage, QrustyClientError>);

/// Sender for log entries routed from the router task.
type LogSender = Arc<Mutex<Option<mpsc::Sender<LogEntry>>>>;

// ---------------------------------------------------------------------------
// LogEntry – matches the server's log frame fields.
// ---------------------------------------------------------------------------

/// A log entry streamed by the server over a WebSocket log subscription.
#[derive(Debug, Clone)]
pub struct LogEntry {
    pub timestamp: String,
    pub level: String,
    pub message: String,
}

// ---------------------------------------------------------------------------
// DeliveredMessage – matches WS-0008 delivery frame fields.
// ---------------------------------------------------------------------------

/// A message delivered by the server over a WebSocket subscription.
#[derive(Debug, Clone)]
pub struct DeliveredMessage {
    pub queue: String,
    pub id: String,
    pub payload: String,
    pub priority: Priority,
    pub created_at: String,
}

impl TryFrom<&Value> for DeliveredMessage {
    type Error = QrustyClientError;

    fn try_from(v: &Value) -> Result<Self, Self::Error> {
        Ok(DeliveredMessage {
            queue: str_field(v, "queue")?,
            id: str_field(v, "id")?,
            payload: str_field(v, "payload")?,
            priority: serde_json::from_value(
                v.get("priority")
                    .cloned()
                    .ok_or_else(|| invalid("missing priority"))?,
            )
            .map_err(|_| invalid("invalid priority value"))?,
            created_at: str_field(v, "created_at")?,
        })
    }
}

// ---------------------------------------------------------------------------
// WsSession – the public API surface (WS-0017).
// ---------------------------------------------------------------------------

/// A persistent WebSocket session to a Qrusty server.
///
/// Create with [`WsSession::connect`] or [`WsSession::connect_with_timeout`].
pub struct WsSession {
    /// Channel for data frames (ack, nack, credit, publish, etc.).
    /// A dedicated writer task drains this into the WS sink, with pong
    /// frames always taking priority (see `sink_writer_task`).
    frame_tx: mpsc::UnboundedSender<TMsg>,
    req_counter: Arc<AtomicU64>,
    pending: PendingMap,
    subscribers: SubMap,
    log_tx: LogSender,
    _router: JoinHandle<()>,
    _delivery_worker: JoinHandle<()>,
    _sink_writer: JoinHandle<()>,
    request_timeout: Duration,
}

impl WsSession {
    // -----------------------------------------------------------------------
    // connect (WS-0017, WS-0026)
    // -----------------------------------------------------------------------

    /// Connect to the Qrusty server at `addr` (e.g. `"ws://localhost:6784"`)
    /// with the default 30-second request timeout.
    pub async fn connect(addr: &str) -> Result<Self, QrustyClientError> {
        Self::connect_with_timeout(addr, Duration::from_secs(30)).await
    }

    /// Connect to the Qrusty server with a custom request timeout (WS-0026).
    pub async fn connect_with_timeout(
        addr: &str,
        request_timeout: Duration,
    ) -> Result<Self, QrustyClientError> {
        let url = format!("{}/ws", addr);
        let (ws, _) = connect_async(&url)
            .await
            .map_err(|e| QrustyClientError::Other(format!("WebSocket connect failed: {}", e)))?;

        let (sink, source) = ws.split();
        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));

        let (delivery_tx, delivery_rx) = mpsc::unbounded_channel();
        let log_tx: LogSender = Arc::new(Mutex::new(None));

        // Data frames (ack, credit, publish, etc.) and pong frames are
        // sent through separate channels.  A single writer task owns the
        // WS sink and uses biased select to always send pongs before
        // data frames, preventing ping timeout under heavy load.
        let (frame_tx, frame_rx) = mpsc::unbounded_channel::<TMsg>();
        let (pong_tx, pong_rx) = mpsc::unbounded_channel::<TMsg>();

        let sink_writer = tokio::spawn(sink_writer_task(sink, frame_rx, pong_rx));

        let router = tokio::spawn(router_task(
            source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        let delivery_worker = tokio::spawn(delivery_task(delivery_rx, Arc::clone(&subscribers)));

        Ok(WsSession {
            frame_tx,
            req_counter: Arc::new(AtomicU64::new(0)),
            pending,
            subscribers,
            log_tx,
            _router: router,
            _delivery_worker: delivery_worker,
            _sink_writer: sink_writer,
            request_timeout,
        })
    }

    // -----------------------------------------------------------------------
    // Internal helpers
    // -----------------------------------------------------------------------

    fn next_req_id(&self) -> String {
        format!("req-{}", self.req_counter.fetch_add(1, Ordering::Relaxed))
    }

    // -----------------------------------------------------------------------
    // send_frame (WS-0025)
    // -----------------------------------------------------------------------

    /// Send a raw JSON frame over the WebSocket without `req_id` correlation.
    ///
    /// This is a fire-and-forget send: no response is awaited.  The frame is
    /// sent as-is.  If you need request-response semantics, use the
    /// higher-level methods (`publish`, `ack`, etc.) instead.
    pub async fn send_frame(&self, frame: Value) -> Result<(), QrustyClientError> {
        self.frame_tx
            .send(TMsg::Text(frame.to_string().into()))
            .map_err(|e| QrustyClientError::Other(format!("send error: {}", e)))
    }

    /// Send a request frame and await its response via a oneshot channel
    /// (WS-0026: with configurable timeout).
    async fn request(&self, mut frame: Value) -> Result<Value, QrustyClientError> {
        let req_id = self.next_req_id();
        frame
            .as_object_mut()
            .unwrap()
            .insert("req_id".to_owned(), Value::String(req_id.clone()));

        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(req_id.clone(), tx);

        // If send fails, remove the pending entry to avoid a leak.
        if let Err(e) = self.send_frame(frame).await {
            self.pending.lock().await.remove(&req_id);
            return Err(e);
        }

        match tokio::time::timeout(self.request_timeout, rx).await {
            Ok(Ok(result)) => result,
            Ok(Err(_)) => Err(QrustyClientError::Other(
                "connection closed before response".into(),
            )),
            Err(_) => {
                self.pending.lock().await.remove(&req_id);
                Err(QrustyClientError::RequestTimeout(self.request_timeout))
            }
        }
    }

    // -----------------------------------------------------------------------
    // publish (WS-0017)
    // -----------------------------------------------------------------------

    /// Publish a message to `queue`.  Returns the assigned message ID.
    pub async fn publish(
        &self,
        queue: &str,
        payload: &str,
        priority: Option<Priority>,
    ) -> Result<String, QrustyClientError> {
        let priority = priority.unwrap_or_default();
        let resp = self
            .request(json!({
                "type": "publish",
                "queue": queue,
                "payload": payload,
                "priority": priority,
            }))
            .await?;
        resp["id"]
            .as_str()
            .map(str::to_owned)
            .ok_or_else(|| invalid("missing id in publish response"))
    }

    // -----------------------------------------------------------------------
    // subscribe (WS-0017, WS-0023)
    // -----------------------------------------------------------------------

    /// Subscribe to `queue` with unlimited delivery (backwards-compatible).
    pub async fn subscribe(
        &self,
        queue: &str,
    ) -> Result<mpsc::Receiver<Result<DeliveredMessage, QrustyClientError>>, QrustyClientError>
    {
        self.subscribe_with_credits(queue, None).await
    }

    /// Subscribe to `queue` with optional credit-based flow control (WS-0023).
    ///
    /// If `credits` is `Some(n)`, the server delivers at most `n` messages
    /// before pausing.  Use [`grant_credits`](Self::grant_credits) to
    /// replenish.  If `credits` is `None`, delivery is unlimited (current
    /// behaviour).
    pub async fn subscribe_with_credits(
        &self,
        queue: &str,
        credits: Option<u64>,
    ) -> Result<mpsc::Receiver<Result<DeliveredMessage, QrustyClientError>>, QrustyClientError>
    {
        self.subscribe_with_options(queue, credits, None).await
    }

    /// Subscribe to `queue` with optional credits and lock timeout.
    ///
    /// `lock_timeout_secs`: how long the server locks each delivered
    /// message before auto-unlocking.  `None` uses the server default
    /// (30 seconds).
    pub async fn subscribe_with_options(
        &self,
        queue: &str,
        credits: Option<u64>,
        lock_timeout_secs: Option<u64>,
    ) -> Result<mpsc::Receiver<Result<DeliveredMessage, QrustyClientError>>, QrustyClientError>
    {
        let (tx, rx) = mpsc::channel::<Result<DeliveredMessage, QrustyClientError>>(256);
        self.subscribers.lock().await.insert(queue.to_owned(), tx);

        let mut frame = json!({"type": "subscribe", "queue": queue});
        if let Some(c) = credits {
            frame
                .as_object_mut()
                .unwrap()
                .insert("credits".to_owned(), json!(c));
        }
        if let Some(t) = lock_timeout_secs {
            frame
                .as_object_mut()
                .unwrap()
                .insert("lock_timeout_secs".to_owned(), json!(t));
        }

        if let Err(e) = self.request(frame).await {
            self.subscribers.lock().await.remove(queue);
            return Err(e);
        }

        Ok(rx)
    }

    // -----------------------------------------------------------------------
    // unsubscribe (WS-0017)
    // -----------------------------------------------------------------------

    /// Unsubscribe from `queue`.
    pub async fn unsubscribe(&self, queue: &str) -> Result<(), QrustyClientError> {
        self.subscribers.lock().await.remove(queue);
        self.request(json!({"type": "unsubscribe", "queue": queue}))
            .await
            .map(|_| ())
    }

    // -----------------------------------------------------------------------
    // ack (WS-0017)
    // -----------------------------------------------------------------------

    /// Acknowledge message `id` on `queue`.
    pub async fn ack(&self, queue: &str, id: &str) -> Result<(), QrustyClientError> {
        self.request(json!({"type": "ack", "queue": queue, "id": id}))
            .await
            .map(|_| ())
    }

    // -----------------------------------------------------------------------
    // nack (WS-0017)
    // -----------------------------------------------------------------------

    /// Negative-acknowledge message `id` on `queue`.
    pub async fn nack(&self, queue: &str, id: &str) -> Result<(), QrustyClientError> {
        self.request(json!({"type": "nack", "queue": queue, "id": id}))
            .await
            .map(|_| ())
    }

    // -----------------------------------------------------------------------
    // ack_noreply / nack_noreply (WS-0024)
    // -----------------------------------------------------------------------

    /// Fire-and-forget acknowledge: sends an ack frame with `no_reply: true`
    /// without waiting for a server response.
    ///
    /// Returns `Ok(())` once the frame has been written to the WebSocket.
    /// Does NOT confirm the server processed the ack successfully.
    pub async fn ack_noreply(&self, queue: &str, id: &str) -> Result<(), QrustyClientError> {
        self.send_frame(json!({
            "type": "ack",
            "queue": queue,
            "id": id,
            "no_reply": true,
        }))
        .await
    }

    /// Fire-and-forget negative-acknowledge: sends a nack frame with
    /// `no_reply: true` without waiting for a server response.
    pub async fn nack_noreply(&self, queue: &str, id: &str) -> Result<(), QrustyClientError> {
        self.send_frame(json!({
            "type": "nack",
            "queue": queue,
            "id": id,
            "no_reply": true,
        }))
        .await
    }

    // -----------------------------------------------------------------------
    // batch_ack (WS-0017)
    // -----------------------------------------------------------------------

    /// Batch-acknowledge multiple messages.  Returns the number acked.
    pub async fn batch_ack(&self, queue: &str, ids: &[&str]) -> Result<usize, QrustyClientError> {
        let resp = self
            .request(json!({"type": "batch-ack", "queue": queue, "ids": ids}))
            .await?;
        resp["acked"]
            .as_u64()
            .map(|n| n as usize)
            .ok_or_else(|| invalid("missing 'acked' field"))
    }

    // -----------------------------------------------------------------------
    // batch_nack (WS-0017)
    // -----------------------------------------------------------------------

    /// Batch-negative-acknowledge multiple messages.
    /// Returns `(unlocked, dropped)`.
    pub async fn batch_nack(
        &self,
        queue: &str,
        ids: &[&str],
    ) -> Result<(usize, usize), QrustyClientError> {
        let resp = self
            .request(json!({"type": "batch-nack", "queue": queue, "ids": ids}))
            .await?;
        let unlocked = resp["unlocked"].as_u64().unwrap_or(0) as usize;
        let dropped = resp["dropped"].as_u64().unwrap_or(0) as usize;
        Ok((unlocked, dropped))
    }

    // -----------------------------------------------------------------------
    // grant_credits (WS-0023)
    // -----------------------------------------------------------------------

    /// Replenish credits for a subscription on `queue` (WS-0023).
    ///
    /// Sends a `credit` frame to the server.  The server adds `credits` to
    /// the subscription's remaining credit count and resumes delivery if it
    /// was paused.
    pub async fn grant_credits(&self, queue: &str, credits: u64) -> Result<(), QrustyClientError> {
        self.request(json!({
            "type": "credit",
            "queue": queue,
            "credits": credits,
        }))
        .await
        .map(|_| ())
    }

    // -----------------------------------------------------------------------
    // subscribe_logs / unsubscribe_logs
    // -----------------------------------------------------------------------

    /// Subscribe to the server's log stream.
    ///
    /// Returns a receiver that yields [`LogEntry`] values as the server
    /// streams them.  The server first replays its buffered history, then
    /// continues with live entries.
    pub async fn subscribe_logs(&self) -> Result<mpsc::Receiver<LogEntry>, QrustyClientError> {
        let (tx, rx) = mpsc::channel::<LogEntry>(256);
        *self.log_tx.lock().await = Some(tx);

        if let Err(e) = self.request(json!({"type": "subscribe-logs"})).await {
            *self.log_tx.lock().await = None;
            return Err(e);
        }

        Ok(rx)
    }

    /// Unsubscribe from the server's log stream.
    pub async fn unsubscribe_logs(&self) -> Result<(), QrustyClientError> {
        self.request(json!({"type": "unsubscribe-logs"}))
            .await
            .map(|_| ())?;
        *self.log_tx.lock().await = None;
        Ok(())
    }

    // -----------------------------------------------------------------------
    // renew (WS-0020)
    // -----------------------------------------------------------------------

    /// Renew (extend) the lock on message `id` in `queue` using the
    /// server default timeout (30 seconds).
    ///
    /// Returns `Ok(())` if the lock was renewed successfully.  Returns an
    /// error if the message is not currently locked by this connection.
    pub async fn renew(&self, queue: &str, id: &str) -> Result<(), QrustyClientError> {
        self.renew_with_timeout(queue, id, None).await
    }

    /// Renew (extend) the lock on message `id` in `queue` with an
    /// optional explicit lock timeout in seconds.
    ///
    /// If `lock_timeout_secs` is `None`, the server uses its default
    /// (30 seconds).
    pub async fn renew_with_timeout(
        &self,
        queue: &str,
        id: &str,
        lock_timeout_secs: Option<u64>,
    ) -> Result<(), QrustyClientError> {
        let mut frame = json!({"type": "renew", "queue": queue, "id": id});
        if let Some(t) = lock_timeout_secs {
            frame
                .as_object_mut()
                .unwrap()
                .insert("lock_timeout_secs".to_owned(), json!(t));
        }
        self.request(frame).await.map(|_| ())
    }

    // -----------------------------------------------------------------------
    // close (WS-0016)
    // -----------------------------------------------------------------------

    /// Close the connection gracefully (WS-0016).
    pub async fn close(self) -> Result<(), QrustyClientError> {
        self.frame_tx
            .send(TMsg::Close(None))
            .map_err(|e| QrustyClientError::Other(format!("close error: {}", e)))
    }
}

// ---------------------------------------------------------------------------
// Demultiplexing router task (WS-0022)
// ---------------------------------------------------------------------------

/// Reads all inbound frames from the WebSocket stream and routes them to
/// the correct destination without holding any long-lived lock:
///
/// - `Ping` frames → forwarded to `pong_tx` for priority reply via the
///   `sink_writer_task`, which always sends pongs before data frames.
/// - `Close` frames → break out of the loop to trigger cleanup.
/// - Frames with a `req_id` → delivered to the matching oneshot sender
///   registered in `pending` by `request()`.
/// - Frames with `"type": "deliver"` → forwarded to the delivery worker
///   via an unbounded channel (never blocks).
/// - All other frames are discarded (e.g. unsolicited server messages).
///
/// When the stream closes, all pending oneshot senders are dropped so that
/// any awaiting `request()` calls receive an error immediately.
// Implements: WS-0022, WS-0028
async fn router_task<S>(
    mut source: S,
    pong_tx: mpsc::UnboundedSender<TMsg>,
    pending: PendingMap,
    subscribers: SubMap,
    delivery_tx: mpsc::UnboundedSender<DeliveryEnvelope>,
    log_tx: LogSender,
) where
    S: futures_util::Stream<Item = Result<TMsg, tokio_tungstenite::tungstenite::Error>> + Unpin,
{
    while let Some(Ok(msg)) = source.next().await {
        let text = match msg {
            TMsg::Text(t) => t,
            TMsg::Ping(data) => {
                let _ = pong_tx.send(TMsg::Pong(data));
                continue;
            }
            TMsg::Close(_) => break,
            _ => continue,
        };
        let frame: Value = match serde_json::from_str(&text) {
            Ok(v) => v,
            Err(_) => continue,
        };

        // Route request-response frames by req_id.
        if let Some(req_id) = frame.get("req_id").and_then(|v| v.as_str()) {
            if let Some(tx) = pending.lock().await.remove(req_id) {
                let result = if frame["type"] == "error" {
                    Err(QrustyClientError::Other(format!(
                        "server error {}: {}",
                        frame["code"].as_str().unwrap_or("?"),
                        frame["message"].as_str().unwrap_or("?"),
                    )))
                } else {
                    Ok(frame)
                };
                let _ = tx.send(result);
            }
            continue;
        }

        // Route deliver frames to the delivery worker (non-blocking).
        if frame["type"] == "deliver" {
            if let Some(queue) = frame["queue"].as_str().map(str::to_owned) {
                let msg = DeliveredMessage::try_from(&frame);
                let _ = delivery_tx.send((queue, msg));
            }
            continue;
        }

        // Route log frames to the log subscriber (non-blocking).
        if frame["type"] == "log" {
            let entry = LogEntry {
                timestamp: frame["timestamp"].as_str().unwrap_or("").to_owned(),
                level: frame["level"].as_str().unwrap_or("").to_owned(),
                message: frame["message"].as_str().unwrap_or("").to_owned(),
            };
            if let Some(tx) = log_tx.lock().await.as_ref() {
                let _ = tx.try_send(entry);
            }
        }
    }

    // Connection closed — wake all pending requests with an error.
    let mut map = pending.lock().await;
    for (_, tx) in map.drain() {
        let _ = tx.send(Err(QrustyClientError::Other("connection closed".into())));
    }

    // Drop all subscriber senders so that any rx.recv() callers unblock
    // instead of waiting forever.
    subscribers.lock().await.clear();
}

// ---------------------------------------------------------------------------
// Delivery worker task (WS-0022)
// ---------------------------------------------------------------------------

/// Reads delivery envelopes from the unbounded channel and forwards them to
/// the appropriate subscriber's bounded mpsc channel using non-blocking
/// `try_send`.  Delivery frames that arrive when the subscriber channel is
/// full are dropped with a warning — the server will redeliver when the
/// message lock expires.
async fn delivery_task(
    mut delivery_rx: mpsc::UnboundedReceiver<DeliveryEnvelope>,
    subscribers: SubMap,
) {
    while let Some((queue, msg)) = delivery_rx.recv().await {
        let maybe_tx = subscribers.lock().await.get(&queue).cloned();
        if let Some(tx) = maybe_tx {
            match tx.try_send(msg) {
                Ok(()) => {}
                Err(mpsc::error::TrySendError::Full(_)) => {
                    log::warn!(
                        "subscriber channel full for queue '{}'; dropping delivery frame",
                        queue,
                    );
                }
                Err(mpsc::error::TrySendError::Closed(_)) => {
                    subscribers.lock().await.remove(&queue);
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Sink writer task — single owner of the WS sink with pong priority
// ---------------------------------------------------------------------------

/// Owns the WS sink and drains two channels: a high-priority pong channel
/// and a regular data-frame channel.  Uses biased `select!` so pong frames
/// are always sent before data frames, preventing ping timeout under heavy
/// outbound traffic.
async fn sink_writer_task<K>(
    mut sink: K,
    mut frame_rx: mpsc::UnboundedReceiver<TMsg>,
    mut pong_rx: mpsc::UnboundedReceiver<TMsg>,
) where
    K: futures_util::Sink<TMsg, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
{
    loop {
        tokio::select! {
            biased;

            // Pong frames always take priority.
            msg = pong_rx.recv() => {
                match msg {
                    Some(pong) => { let _ = SinkExt::send(&mut sink, pong).await; }
                    None => break, // pong channel closed
                }
            }

            // Data frames (ack, nack, credit, publish, close, etc.).
            msg = frame_rx.recv() => {
                match msg {
                    Some(frame) => { let _ = SinkExt::send(&mut sink, frame).await; }
                    None => break, // data channel closed
                }
            }
        }
    }

    // Attempt graceful close.
    let _ = SinkExt::close(&mut sink).await;
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn str_field(v: &Value, field: &str) -> Result<String, QrustyClientError> {
    v.get(field)
        .and_then(|f| f.as_str())
        .map(str::to_owned)
        .ok_or_else(|| invalid(&format!("missing field '{}'", field)))
}

fn invalid(msg: &str) -> QrustyClientError {
    QrustyClientError::InvalidResponse(msg.to_owned())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::stream::SplitSink;
    use tokio_tungstenite::tungstenite::protocol::Role;

    /// Build an in-memory WebSocket pair without a real network connection.
    /// Returns `(client_source, client_sink, pong_tx, pong_rx, server_ws)`.
    /// - `router_task` reads from `client_source` and pushes pongs into `pong_tx`.
    /// - `sink_writer_task` drains `pong_rx` and `frame_rx` into `client_sink`.
    /// - The test controls `server_ws`.
    async fn in_memory_ws_pair() -> (
        futures_util::stream::SplitStream<
            tokio_tungstenite::WebSocketStream<tokio::io::DuplexStream>,
        >,
        SplitSink<tokio_tungstenite::WebSocketStream<tokio::io::DuplexStream>, TMsg>,
        mpsc::UnboundedSender<TMsg>,
        mpsc::UnboundedReceiver<TMsg>,
        tokio_tungstenite::WebSocketStream<tokio::io::DuplexStream>,
    ) {
        let (server_io, client_io) = tokio::io::duplex(65_536);
        let server_ws =
            tokio_tungstenite::WebSocketStream::from_raw_socket(server_io, Role::Server, None)
                .await;
        let client_ws =
            tokio_tungstenite::WebSocketStream::from_raw_socket(client_io, Role::Client, None)
                .await;
        let (client_sink, client_source) = client_ws.split();
        let (pong_tx, pong_rx) = mpsc::unbounded_channel();
        (client_source, client_sink, pong_tx, pong_rx, server_ws)
    }

    /// Dropping the server side of an in-memory WebSocket causes the client
    /// source to see EOF.  `router_task` must then:
    ///   - send a "connection closed" error to every pending oneshot, and
    ///   - drop all subscriber mpsc senders so that `rx.recv()` returns None.
    // Verifies: WS-0017
    #[tokio::test]
    async fn router_close_wakes_pending_and_clears_subscribers() {
        let (client_source, _client_sink, pong_tx, _pong_rx, server_ws) = in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();

        // Register one pending request so we can observe it being woken.
        let (tx, rx) = oneshot::channel();
        pending.lock().await.insert("req-0".to_owned(), tx);

        // Register one subscriber so we can observe the channel being closed.
        let (sub_tx, mut sub_rx) = mpsc::channel(8);
        subscribers.lock().await.insert("orders".to_owned(), sub_tx);

        // Spawn the router then immediately close the server side.
        let log_tx: LogSender = Arc::new(Mutex::new(None));
        let task = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));
        drop(server_ws);
        task.await.unwrap();

        // Pending oneshot must carry a "connection closed" error.
        let result = rx.await.expect("oneshot was not sent");
        assert!(result.is_err(), "expected Err, got Ok");

        // Subscriber channel must be closed (sender was dropped by clear()).
        assert!(
            sub_rx.recv().await.is_none(),
            "subscriber channel should be closed after connection close"
        );

        // Both registries must be empty.
        assert!(pending.lock().await.is_empty());
        assert!(subscribers.lock().await.is_empty());
    }

    /// WS-0022: router_task does not block when subscriber channel is full.
    /// The delivery worker uses try_send, so request-response correlation
    /// continues even if a subscriber is backpressured.
    // Verifies: WS-0022
    #[tokio::test]
    async fn router_does_not_block_on_full_subscriber_channel() {
        let (client_source, _client_sink, pong_tx, _pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, delivery_rx) = mpsc::unbounded_channel();

        // Subscriber channel with buffer=1 (tiny).
        let (sub_tx, _sub_rx) = mpsc::channel(1);
        subscribers.lock().await.insert("q".to_owned(), sub_tx);

        // Pending request we will use to prove the router is not stuck.
        let (req_tx, req_rx) = oneshot::channel();
        pending.lock().await.insert("req-1".to_owned(), req_tx);

        let log_tx: LogSender = Arc::new(Mutex::new(None));
        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));
        let delivery_handle = tokio::spawn(delivery_task(delivery_rx, Arc::clone(&subscribers)));

        // Send 5 deliver frames (channel buffer=1, so most will be dropped).
        for i in 0..5 {
            let frame = json!({
                "type": "deliver",
                "queue": "q",
                "id": format!("msg-{}", i),
                "payload": "data",
                "priority": 0,
                "created_at": "2026-01-01T00:00:00Z",
            });
            server_ws
                .send(TMsg::Text(frame.to_string().into()))
                .await
                .unwrap();
        }

        // Send a req_id response.  If the router was blocked by try_send
        // this would never be delivered.
        let resp_frame = json!({"req_id": "req-1", "type": "ok"});
        server_ws
            .send(TMsg::Text(resp_frame.to_string().into()))
            .await
            .unwrap();

        // The pending request must resolve (proves router was not blocked).
        let result = tokio::time::timeout(Duration::from_secs(2), req_rx)
            .await
            .expect("timed out waiting for req-1 response")
            .expect("oneshot recv error");
        assert!(result.is_ok());

        drop(server_ws);
        let _ = router_handle.await;
        let _ = delivery_handle.await;
    }

    /// WS-0026: request times out if no response arrives.
    // Verifies: WS-0026
    #[tokio::test]
    async fn request_timeout_cleans_up_pending() {
        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));

        let (tx, rx) = oneshot::channel::<Result<Value, QrustyClientError>>();
        pending.lock().await.insert("req-timeout".to_owned(), tx);

        let timeout_dur = Duration::from_millis(50);
        let result = tokio::time::timeout(timeout_dur, rx).await;
        assert!(result.is_err(), "should have timed out");

        // Caller would remove the entry on timeout.
        pending.lock().await.remove("req-timeout");
        assert!(pending.lock().await.is_empty());
    }

    /// Router routes "type": "log" frames to the log channel.
    #[tokio::test]
    async fn router_routes_log_frames_to_log_channel() {
        let (client_source, _client_sink, pong_tx, _pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();
        let (log_tx_inner, mut log_rx) = mpsc::channel::<LogEntry>(256);
        let log_tx: LogSender = Arc::new(Mutex::new(Some(log_tx_inner)));

        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        // Send a log frame from the server.
        let log_frame = json!({
            "type": "log",
            "timestamp": "2026-03-10T00:00:00Z",
            "level": "INFO",
            "message": "test log message",
        });
        server_ws
            .send(TMsg::Text(log_frame.to_string().into()))
            .await
            .unwrap();

        // The log entry should arrive in the channel.
        let entry = tokio::time::timeout(Duration::from_secs(2), log_rx.recv())
            .await
            .expect("timed out waiting for log entry")
            .expect("log channel closed");

        assert_eq!(entry.timestamp, "2026-03-10T00:00:00Z");
        assert_eq!(entry.level, "INFO");
        assert_eq!(entry.message, "test log message");

        drop(server_ws);
        let _ = router_handle.await;
    }

    /// Log frames are silently dropped when no log subscriber is registered.
    #[tokio::test]
    async fn router_drops_log_frames_when_no_subscriber() {
        let (client_source, _client_sink, pong_tx, _pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();
        let log_tx: LogSender = Arc::new(Mutex::new(None)); // no subscriber

        // Register a pending request to prove the router processes frames.
        let (req_tx, req_rx) = oneshot::channel();
        pending.lock().await.insert("req-1".to_owned(), req_tx);

        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        // Send a log frame (should be silently dropped).
        let log_frame = json!({
            "type": "log",
            "timestamp": "2026-03-10T00:00:00Z",
            "level": "WARN",
            "message": "dropped",
        });
        server_ws
            .send(TMsg::Text(log_frame.to_string().into()))
            .await
            .unwrap();

        // Send a request-response to prove the router is still working.
        let resp_frame = json!({"req_id": "req-1", "type": "ok"});
        server_ws
            .send(TMsg::Text(resp_frame.to_string().into()))
            .await
            .unwrap();

        let result = tokio::time::timeout(Duration::from_secs(2), req_rx)
            .await
            .expect("timed out waiting for req-1")
            .expect("oneshot recv error");
        assert!(result.is_ok());

        drop(server_ws);
        let _ = router_handle.await;
    }

    /// Router replies to Ping frames with Pong to keep the connection alive.
    // Verifies: WS-0028
    #[tokio::test]
    async fn router_replies_pong_to_ping() {
        let (client_source, client_sink, pong_tx, pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();
        let (frame_tx, frame_rx) = mpsc::unbounded_channel();
        let log_tx: LogSender = Arc::new(Mutex::new(None));

        // Start the sink writer so pongs reach the server.
        let writer_handle = tokio::spawn(sink_writer_task(client_sink, frame_rx, pong_rx));

        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        // Send a Ping from the server.
        server_ws
            .send(TMsg::Ping(vec![1, 2, 3].into()))
            .await
            .unwrap();

        // Server should receive a Pong with the same payload.
        let reply = tokio::time::timeout(Duration::from_secs(2), server_ws.next())
            .await
            .expect("timed out waiting for Pong")
            .expect("stream ended")
            .expect("frame error");
        assert_eq!(reply, TMsg::Pong(vec![1, 2, 3].into()));

        drop(server_ws);
        drop(frame_tx);
        let _ = router_handle.await;
        let _ = writer_handle.await;
    }

    /// Pong frames are sent before data frames even when data frames are
    /// queued first.  The sink_writer_task uses biased select to prioritize
    /// pong over data.
    // Verifies: WS-0028
    #[tokio::test]
    async fn pong_prioritized_over_data_frames() {
        let (client_source, client_sink, pong_tx, pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();
        let (frame_tx, frame_rx) = mpsc::unbounded_channel();
        let log_tx: LogSender = Arc::new(Mutex::new(None));

        // Queue a data frame BEFORE the sink writer starts.
        frame_tx
            .send(TMsg::Text(r#"{"type":"noop"}"#.into()))
            .unwrap();

        // Queue a pong frame (simulating what router_task does on Ping).
        pong_tx.send(TMsg::Pong(vec![42].into())).unwrap();

        // Now start the sink writer — biased select should send pong first.
        let writer_handle = tokio::spawn(sink_writer_task(client_sink, frame_rx, pong_rx));

        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        // First frame received by server should be the Pong (priority).
        let first = tokio::time::timeout(Duration::from_secs(2), server_ws.next())
            .await
            .expect("timed out")
            .expect("stream ended")
            .expect("frame error");
        assert_eq!(first, TMsg::Pong(vec![42].into()));

        drop(server_ws);
        drop(frame_tx);
        let _ = router_handle.await;
        let _ = writer_handle.await;
    }

    /// Router breaks out of its loop on Close frames and cleans up.
    // Verifies: WS-0028
    #[tokio::test]
    async fn router_handles_close_frame() {
        let (client_source, _client_sink, pong_tx, _pong_rx, mut server_ws) =
            in_memory_ws_pair().await;

        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
        let subscribers: SubMap = Arc::new(Mutex::new(HashMap::new()));
        let (delivery_tx, _delivery_rx) = mpsc::unbounded_channel();
        let log_tx: LogSender = Arc::new(Mutex::new(None));

        // Register a subscriber to verify it gets cleaned up.
        let (sub_tx, mut sub_rx) = mpsc::channel(8);
        subscribers.lock().await.insert("q".to_owned(), sub_tx);

        let router_handle = tokio::spawn(router_task(
            client_source,
            pong_tx,
            Arc::clone(&pending),
            Arc::clone(&subscribers),
            delivery_tx,
            Arc::clone(&log_tx),
        ));

        // Send a Close frame from the server.
        server_ws.send(TMsg::Close(None)).await.unwrap();

        // Router should exit and clean up subscribers.
        let _ = router_handle.await;
        assert!(
            sub_rx.recv().await.is_none(),
            "subscriber channel should be closed after Close frame"
        );
        assert!(subscribers.lock().await.is_empty());
    }
}
