qrusty 0.20.9

// src/ws.rs
//
// WebSocket handler for the Qrusty priority queue server.
//
// Implements: WS-0001, WS-0002, WS-0003, WS-0004, WS-0005, WS-0006, WS-0007, WS-0008, WS-0009, WS-0010, WS-0011, WS-0012, WS-0013, WS-0014, WS-0015, WS-0016, WS-0020, WS-0023, WS-0024, WS-0027
//
// Architecture
// ============
//
// Each connection spawns three categories of tasks:
//
//   1. Main loop  (`handle_connection`)
//      Multiplexes between:
//        a. Incoming WS frames from the client
//        b. Deliveries from per-subscription poller tasks (via `deliver_rx`)
//        c. Log entries from the log-buffer subscriber task (via `log_rx`)
//        d. Server-side ping ticker (WS-0015)
//      The main loop places outbound frames into `outbound_tx` and never
//      writes directly to the WS socket (WS-0027).
//
//   2. Sender task
//      Drains `outbound_rx` and writes frames to the WS socket.  Network
//      back-pressure is therefore isolated to this task; the main loop
//      continues processing inbound frames and enforcing ping timeouts even
//      when the socket write stalls (WS-0027).
//
//   3. Poller tasks  (one per active subscription)
//      Loop calling `storage.pop()` and forward messages to the main loop
//      via `deliver_tx` (WS-0007, WS-0023).

use crate::api::{QueueRateTracker, StorageApi};
use crate::message::Priority;
use axum::extract::ws::{CloseFrame, Message as WsMessage, WebSocket};
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{mpsc, Notify};
use uuid::Uuid;

// ---------------------------------------------------------------------------
// Tuning
// ---------------------------------------------------------------------------

// Environment-variable controlled timings (WS-0015).
fn ping_interval() -> Duration {
    let secs: u64 = std::env::var("WS_PING_INTERVAL_SECS")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(30);
    Duration::from_secs(secs)
}

fn ping_timeout() -> Duration {
    let secs: u64 = std::env::var("WS_PING_TIMEOUT_SECS")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(30);
    Duration::from_secs(secs)
}

/// Capacity of the outbound channel between the main loop and the sender task.
///
/// When this channel fills the main loop drops the frame rather than blocking,
/// preserving liveness of the receive/ping path (WS-0027).
const OUTBOUND_CHANNEL_CAPACITY: usize = 512;

// ---------------------------------------------------------------------------
// Delivery payload – sent from poller tasks to the main connection loop.
// ---------------------------------------------------------------------------

struct Delivery {
    queue: String,
    id: String,
    payload: String,
    priority: Priority,
    created_at: String,
}

impl Delivery {
    fn to_frame(&self) -> WsMessage {
        let v = json!({
            "type": "deliver",
            "queue":      self.queue,
            "id":         self.id,
            "payload":    self.payload,
            "priority":   self.priority,
            "created_at": self.created_at,
        });
        WsMessage::Text(v.to_string().into())
    }
}

// ---------------------------------------------------------------------------
// SubscriptionRegistry – tracks active poller task handles per queue.
// ---------------------------------------------------------------------------

/// Per-subscription state: abort handle, credit counter, and wake signals.
struct SubscriptionState {
    abort_handle: tokio::task::AbortHandle,
    /// Remaining credits.  Negative means unlimited.
    credits: Arc<AtomicI64>,
    /// Notified when credits are replenished (wakes the paused poller).
    credits_notify: Arc<Notify>,
    /// Notified when a message is published to this queue (wakes the
    /// poller from its empty-queue sleep so delivery is immediate).
    message_notify: Arc<Notify>,
}

struct SubscriptionRegistry {
    handles: HashMap<String, SubscriptionState>,
}

impl SubscriptionRegistry {
    fn new() -> Self {
        Self {
            handles: HashMap::new(),
        }
    }

    /// Spawn a poller task for `queue` if not already subscribed (WS-0023).
    ///
    /// `credits`: `Some(n)` for credit-limited, `None` for unlimited.
    /// `lock_timeout_secs`: lock duration for popped messages (default: 30).
    fn subscribe(
        &mut self,
        queue: String,
        storage: Arc<dyn StorageApi>,
        conn_id: String,
        tx: mpsc::Sender<Delivery>,
        credits: Option<u64>,
        lock_timeout_secs: u64,
    ) {
        if self.handles.contains_key(&queue) {
            return; // already subscribed
        }

        let credit_count = Arc::new(AtomicI64::new(
            credits.map(|c| c as i64).unwrap_or(-1), // -1 = unlimited
        ));
        let credits_notify = Arc::new(Notify::new());
        let message_notify = Arc::new(Notify::new());

        let cc = Arc::clone(&credit_count);
        let cn = Arc::clone(&credits_notify);
        let mn = Arc::clone(&message_notify);
        let q = queue.clone();

        let handle = tokio::spawn(async move {
            loop {
                // Check credits before popping (WS-0023).
                let remaining = cc.load(AtomicOrdering::Acquire);
                if remaining == 0 {
                    // No credits left — wait for replenishment.
                    cn.notified().await;
                    continue;
                }

                match storage.pop(&q, &conn_id, lock_timeout_secs).await {
                    Ok(Some(msg)) => {
                        let delivery = Delivery {
                            queue: msg.queue.clone(),
                            id: msg.id.clone(),
                            payload: msg.payload.clone(),
                            priority: msg.priority,
                            created_at: msg.created_at.to_rfc3339(),
                        };
                        if tx.send(delivery).await.is_err() {
                            break; // receiver dropped → connection gone
                        }
                        // Decrement credit if not unlimited (-1).
                        if remaining > 0 {
                            cc.fetch_sub(1, AtomicOrdering::Release);
                        }
                    }
                    Ok(None) => {
                        // No message available — wait for a publish or
                        // time out after 500ms.  The message_notify is
                        // signalled by notify_message_available() when a
                        // message is published to this queue via WS.
                        // HTTP publishes bypass the WS subscription
                        // registry, so the fallback timeout ensures
                        // those messages are still picked up promptly.
                        tokio::time::timeout(Duration::from_millis(500), mn.notified())
                            .await
                            .ok();
                    }
                    Err(_) => {
                        tokio::time::sleep(Duration::from_millis(500)).await;
                    }
                }
            }
        })
        .abort_handle();

        self.handles.insert(
            queue,
            SubscriptionState {
                abort_handle: handle,
                credits: credit_count,
                credits_notify,
                message_notify,
            },
        );
    }

    fn unsubscribe(&mut self, queue: &str) {
        if let Some(state) = self.handles.remove(queue) {
            state.abort_handle.abort();
        }
    }

    /// Add credits to a subscription's counter and wake the poller (WS-0023).
    fn add_credits(&self, queue: &str, n: u64) -> bool {
        if let Some(state) = self.handles.get(queue) {
            let current = state.credits.load(AtomicOrdering::Acquire);
            if current >= 0 {
                state.credits.fetch_add(n as i64, AtomicOrdering::Release);
                state.credits_notify.notify_one();
            }
            // If current < 0 (unlimited), adding credits is a no-op.
            true
        } else {
            false
        }
    }

    /// Wake the poller for `queue` so it checks for new messages immediately.
    /// Called after a message is published to this queue.
    fn notify_message_available(&self, queue: &str) {
        if let Some(state) = self.handles.get(queue) {
            state.message_notify.notify_one();
        }
    }
}

impl Drop for SubscriptionRegistry {
    fn drop(&mut self) {
        for state in self.handles.values() {
            state.abort_handle.abort();
        }
    }
}

// ---------------------------------------------------------------------------
// Main connection handler (called by the Axum WebSocket upgrade).
// ---------------------------------------------------------------------------

// Implements: WS-0021, WS-0027
pub(crate) async fn handle_connection(
    ws: WebSocket,
    storage: Arc<dyn StorageApi>,
    rate_tracker: Arc<QueueRateTracker>,
    log_buffer: Option<crate::log_buffer::LogBuffer>,
    memory_pressure: Arc<std::sync::atomic::AtomicBool>,
) {
    let conn_id = Uuid::new_v4().to_string();
    tracing::info!(conn_id = %conn_id, "ws: connection opened");
    let conn_started = Instant::now();
    let (ws_sender, mut receiver) = ws.split();

    // -----------------------------------------------------------------------
    // Outbound sender task (WS-0027)
    //
    // The main loop places frames onto `outbound_tx`.  A dedicated task drains
    // the channel and writes to the socket.  This decouples slow network writes
    // from the receive/ping path: the main loop never blocks on a socket write.
    // -----------------------------------------------------------------------
    let (outbound_tx, mut outbound_rx) = mpsc::channel::<WsMessage>(OUTBOUND_CHANNEL_CAPACITY);

    let sender_conn_id = conn_id.clone();
    let sender_handle = tokio::spawn(async move {
        let mut sender = ws_sender;
        let mut frames_sent: u64 = 0;
        while let Some(frame) = outbound_rx.recv().await {
            if let Err(e) = sender.send(frame).await {
                tracing::warn!(
                    conn_id = %sender_conn_id,
                    frames_sent,
                    "ws: outbound sender: socket write error, exiting: {e}"
                );
                break;
            }
            frames_sent += 1;
        }
        // Drain and flush on exit.
        if let Err(e) = sender.flush().await {
            tracing::debug!(
                conn_id = %sender_conn_id,
                "ws: outbound sender: flush on exit failed: {e}"
            );
        }
    });

    // -----------------------------------------------------------------------
    // Helper: enqueue an outbound frame without blocking (WS-0027).
    //
    // Drops the frame if the outbound channel is full so the main loop is
    // never stalled by back-pressure.  A dropped delivery frame means the
    // client never sees that message and never grants a credit for it, so
    // we always log dropped frames at WARN: an unnoticed silent drop is
    // the shape of a delivery-stall bug (messages vanish from in-flight
    // without anyone noticing).
    // -----------------------------------------------------------------------
    macro_rules! send_frame {
        ($frame:expr) => {
            if let Err(__send_err) = outbound_tx.try_send($frame) {
                use tokio::sync::mpsc::error::TrySendError;
                match __send_err {
                    TrySendError::Full(__dropped) => {
                        let __kind = match &__dropped {
                            WsMessage::Text(_) => "text",
                            WsMessage::Binary(_) => "binary",
                            WsMessage::Ping(_) => "ping",
                            WsMessage::Pong(_) => "pong",
                            WsMessage::Close(_) => "close",
                        };
                        tracing::warn!(
                            conn_id = %conn_id,
                            frame_type = __kind,
                            capacity = OUTBOUND_CHANNEL_CAPACITY,
                            "ws: outbound channel full, dropped frame"
                        );
                    }
                    TrySendError::Closed(_) => {
                        // Sender task exited; main loop will exit naturally
                        // via the receiver.next() arm on the next inbound
                        // frame (or stay alive doing nothing useful).
                        tracing::debug!(
                            conn_id = %conn_id,
                            "ws: outbound channel closed, frame not sent"
                        );
                    }
                }
            }
        };
    }

    // Channel for poller → main-loop message delivery.
    let (deliver_tx, mut deliver_rx) = mpsc::channel::<Delivery>(256);

    // Channel for log entries → main-loop forwarding.
    let (log_tx, mut log_rx) = mpsc::channel::<crate::log_buffer::LogEntry>(256);

    let ping_iv = ping_interval();
    let ping_to = ping_timeout();
    let mut ping_ticker = tokio::time::interval(ping_iv);
    ping_ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
    // Skip the immediate first tick (fires at t=0).
    ping_ticker.tick().await;

    let mut last_pong = Instant::now();
    let mut subscriptions = SubscriptionRegistry::new();
    let mut log_task: Option<tokio::task::JoinHandle<()>> = None;

    loop {
        tokio::select! {
            // ----------------------------------------------------------------
            // Incoming frame from the WebSocket client
            // ----------------------------------------------------------------
            msg = receiver.next() => {
                let ws_msg = match msg {
                    Some(Ok(m)) => m,
                    Some(Err(e)) => {
                        tracing::warn!(
                            conn_id = %conn_id,
                            uptime_secs = conn_started.elapsed().as_secs_f64(),
                            "ws: receiver error, closing: {e}"
                        );
                        break;
                    }
                    None => {
                        tracing::info!(
                            conn_id = %conn_id,
                            uptime_secs = conn_started.elapsed().as_secs_f64(),
                            "ws: receiver stream ended (EOF)"
                        );
                        break;
                    }
                };

                match ws_msg {
                    WsMessage::Binary(_) => {
                        tracing::warn!(
                            conn_id = %conn_id,
                            "ws: received binary frame, closing with 1003"
                        );
                        // WS-0002: binary frames → 1003 close
                        send_frame!(WsMessage::Close(Some(CloseFrame {
                            code: axum::extract::ws::close_code::UNSUPPORTED,
                            reason: "binary frames not supported".into(),
                        })));
                        break;
                    }
                    WsMessage::Text(text) => {
                        // Parse JSON.
                        let frame: Value = match serde_json::from_str(&text) {
                            Ok(v) => v,
                            Err(e) => {
                                tracing::warn!(
                                    conn_id = %conn_id,
                                    "ws: invalid JSON, closing with 1007: {e}"
                                );
                                // WS-0002: invalid JSON → 1007 close
                                send_frame!(WsMessage::Close(Some(CloseFrame {
                                    code: axum::extract::ws::close_code::INVALID,
                                    reason: "invalid JSON".into(),
                                })));
                                break;
                            }
                        };
                        let req_id = frame.get("req_id").and_then(|v| v.as_str()).map(str::to_owned);
                        let no_reply = frame.get("no_reply").and_then(|v| v.as_bool()).unwrap_or(false);

                        let frame_type = frame.get("type").and_then(|v| v.as_str()).unwrap_or("");

                        let out = match frame_type {
                            "subscribe-logs" => {
                                handle_subscribe_logs(&log_buffer, &mut log_task, &log_tx).await
                            }
                            "unsubscribe-logs" => {
                                handle_unsubscribe_logs(&mut log_task)
                            }
                            _ => {
                                dispatch(
                                    &frame,
                                    &storage,
                                    &rate_tracker,
                                    &mut subscriptions,
                                    &conn_id,
                                    &deliver_tx,
                                    &memory_pressure,
                                )
                                .await
                            }
                        };

                        // WS-0024: suppress response when no_reply is set.
                        if !no_reply {
                            let out = attach_req_id(out, req_id);
                            send_frame!(WsMessage::Text(out.to_string().into()));
                        }
                    }
                    WsMessage::Ping(data) => {
                        // Echo pong (RFC 6455).
                        send_frame!(WsMessage::Pong(data));
                    }
                    WsMessage::Pong(_) => {
                        last_pong = Instant::now();
                    }
                    WsMessage::Close(cf) => {
                        match cf {
                            Some(cf) => tracing::info!(
                                conn_id = %conn_id,
                                uptime_secs = conn_started.elapsed().as_secs_f64(),
                                "ws: client Close frame: code={} reason='{}'",
                                cf.code,
                                cf.reason
                            ),
                            None => tracing::info!(
                                conn_id = %conn_id,
                                uptime_secs = conn_started.elapsed().as_secs_f64(),
                                "ws: client Close frame (no payload)"
                            ),
                        }
                        break;
                    }
                }
            }

            // ----------------------------------------------------------------
            // Delivery frame from a subscription poller
            // ----------------------------------------------------------------
            Some(delivery) = deliver_rx.recv() => {
                send_frame!(delivery.to_frame());
            }

            // ----------------------------------------------------------------
            // Log entry from the log buffer subscriber
            // ----------------------------------------------------------------
            Some(entry) = log_rx.recv() => {
                let frame = json!({
                    "type": "log",
                    "timestamp": entry.timestamp,
                    "level": entry.level,
                    "message": entry.message,
                });
                send_frame!(WsMessage::Text(frame.to_string().into()));
            }

            // ----------------------------------------------------------------
            // Server-side ping (WS-0015)
            // ----------------------------------------------------------------
            _ = ping_ticker.tick() => {
                // last_pong is set to Instant::now() on connect and on every
                // received pong.  After sending a ping we do NOT reset it —
                // we want to detect the case where the client stops answering.
                // After two consecutive unanswered ticks the elapsed time
                // exceeds ping_iv + ping_to and we close with 1001 (AWAY).
                if last_pong.elapsed() > ping_iv + ping_to {
                    tracing::warn!(
                        conn_id = %conn_id,
                        uptime_secs = conn_started.elapsed().as_secs_f64(),
                        last_pong_secs_ago = last_pong.elapsed().as_secs_f64(),
                        "ws: ping timeout, closing with 1001"
                    );
                    send_frame!(WsMessage::Close(Some(CloseFrame {
                        code: axum::extract::ws::close_code::AWAY,
                        reason: "ping timeout".into(),
                    })));
                    break;
                }
                tracing::debug!(
                    conn_id = %conn_id,
                    last_pong_secs_ago = last_pong.elapsed().as_secs_f64(),
                    "ws: sending ping"
                );
                send_frame!(WsMessage::Ping(vec![].into()));
            }
        }
    }

    let subscription_count = subscriptions.handles.len();
    tracing::info!(
        conn_id = %conn_id,
        uptime_secs = conn_started.elapsed().as_secs_f64(),
        subscriptions = subscription_count,
        "ws: connection main loop exited, starting shutdown"
    );

    // Drop the outbound sender so the sender task shuts down cleanly.
    drop(outbound_tx);

    // SubscriptionRegistry::drop() aborts all poller tasks (WS-0016).
    // Abort any active log subscription task.
    if let Some(task) = log_task.take() {
        task.abort();
    }

    // Wait for the sender task to drain and flush.
    let _ = sender_handle.await;

    tracing::info!(
        conn_id = %conn_id,
        total_uptime_secs = conn_started.elapsed().as_secs_f64(),
        "ws: connection closed"
    );
}

// ---------------------------------------------------------------------------
// Frame dispatch
// ---------------------------------------------------------------------------

async fn dispatch(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    subs: &mut SubscriptionRegistry,
    conn_id: &str,
    deliver_tx: &mpsc::Sender<Delivery>,
    memory_pressure: &std::sync::atomic::AtomicBool,
) -> Value {
    let frame_type = match frame.get("type").and_then(|v| v.as_str()) {
        Some(t) => t,
        None => {
            return error_frame("invalid_request", "missing 'type' field");
        }
    };

    match frame_type {
        "publish" => handle_publish(frame, storage, rate_tracker, memory_pressure, subs).await,
        "subscribe" => handle_subscribe(frame, storage, subs, conn_id, deliver_tx).await,
        "unsubscribe" => handle_unsubscribe(frame, subs),
        "ack" => handle_ack(frame, storage, rate_tracker, conn_id).await,
        "nack" => handle_nack(frame, storage, rate_tracker, conn_id).await,
        "batch-ack" => handle_batch_ack(frame, storage, rate_tracker, conn_id).await,
        "batch-nack" => handle_batch_nack(frame, storage, rate_tracker, conn_id).await,
        "renew" => handle_renew(frame, storage, conn_id).await,
        "credit" => handle_credit(frame, subs),
        other => error_frame("unknown_type", &format!("unknown frame type '{}'", other)),
    }
}

// ---------------------------------------------------------------------------
// Individual operation handlers
// ---------------------------------------------------------------------------

/// WS-0006, WS-0021: publish a message.
async fn handle_publish(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    memory_pressure: &std::sync::atomic::AtomicBool,
    subs: &SubscriptionRegistry,
) -> Value {
    // Load shedding: reject WS publishes under memory pressure (SYS-0022).
    if memory_pressure.load(std::sync::atomic::Ordering::Relaxed) {
        return error_frame(
            "memory_pressure",
            "Server is under memory pressure. Try again later.",
        );
    }

    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q.to_owned(),
        None => return error_frame("invalid_request", "publish requires 'queue' field"),
    };
    let payload = match frame.get("payload").and_then(|v| v.as_str()) {
        Some(p) => p.to_owned(),
        None => return error_frame("invalid_request", "publish requires 'payload' field"),
    };
    let priority = match frame.get("priority") {
        Some(serde_json::Value::String(s)) => Priority::Text(s.clone()),
        Some(v) => Priority::Numeric(v.as_u64().unwrap_or(0)),
        None => Priority::Numeric(0),
    };
    let max_retries = frame
        .get("max_retries")
        .and_then(|v| v.as_u64())
        .map(|v| v as u32);

    // Verify the queue exists before pushing so we can return a proper error.
    match storage.queue_exists(&queue).await {
        Ok(true) => {}
        Ok(false) => {
            return error_frame("queue_not_found", &format!("queue '{}' not found", queue))
        }
        Err(e) => return error_frame("storage_error", &e.to_string()),
    }

    let msg = crate::message::Message {
        id: Uuid::new_v4().to_string(),
        queue: queue.clone(),
        priority,
        payload,
        created_at: chrono::Utc::now(),
        locked_until: None,
        locked_by: None,
        retry_count: 0,
        max_retries: max_retries.unwrap_or(3),
        payload_ref: None,
        payload_hash: None,
    };

    match storage.push(msg).await {
        Ok(id) => {
            rate_tracker.record_publish(&queue);
            // Wake any subscriber poller sleeping on an empty queue.
            subs.notify_message_available(&queue);
            json!({"type": "ok", "id": id})
        }
        Err(e) => error_frame("publish_failed", &e.to_string()),
    }
}

/// WS-0007, WS-0023: subscribe to a queue.
async fn handle_subscribe(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    subs: &mut SubscriptionRegistry,
    conn_id: &str,
    deliver_tx: &mpsc::Sender<Delivery>,
) -> Value {
    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q.to_owned(),
        None => return error_frame("invalid_request", "subscribe requires 'queue' field"),
    };
    let credits = frame.get("credits").and_then(|v| v.as_u64());
    let lock_timeout_secs = frame
        .get("lock_timeout_secs")
        .and_then(|v| v.as_u64())
        .unwrap_or(30);

    match storage.queue_exists(&queue).await {
        Ok(true) => {}
        Ok(false) => {
            return error_frame("queue_not_found", &format!("queue '{}' not found", queue))
        }
        Err(e) => return error_frame("storage_error", &e.to_string()),
    }

    subs.subscribe(
        queue,
        storage.clone(),
        conn_id.to_owned(),
        deliver_tx.clone(),
        credits,
        lock_timeout_secs,
    );
    json!({"type": "ok"})
}

/// WS-0009: unsubscribe from a queue.
fn handle_unsubscribe(frame: &Value, subs: &mut SubscriptionRegistry) -> Value {
    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q,
        None => return error_frame("invalid_request", "unsubscribe requires 'queue' field"),
    };
    subs.unsubscribe(queue);
    json!({"type": "ok"})
}

/// WS-0011, WS-0021: acknowledge a message.
async fn handle_ack(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    conn_id: &str,
) -> Value {
    let (queue, id) = match required_queue_id(frame) {
        Ok(v) => v,
        Err(e) => return e,
    };
    match storage.ack(&queue, &id, conn_id).await {
        Ok(true) => {
            rate_tracker.record_ack(&queue);
            json!({"type": "ok"})
        }
        Ok(false) => json!({"type": "ok"}),
        Err(e) => error_frame("ack_failed", &e.to_string()),
    }
}

/// WS-0012, WS-0021: negative-acknowledge a message.
async fn handle_nack(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    conn_id: &str,
) -> Value {
    let (queue, id) = match required_queue_id(frame) {
        Ok(v) => v,
        Err(e) => return e,
    };
    match storage.nack(&queue, &id, conn_id).await {
        Ok(true) => {
            rate_tracker.record_nack(&queue);
            json!({"type": "ok"})
        }
        Ok(false) => json!({"type": "ok"}),
        Err(e) => error_frame("nack_failed", &e.to_string()),
    }
}

/// WS-0013, WS-0021: batch acknowledge.
async fn handle_batch_ack(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    conn_id: &str,
) -> Value {
    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q.to_owned(),
        None => return error_frame("invalid_request", "batch-ack requires 'queue' field"),
    };
    let ids: Vec<String> = match frame.get("ids").and_then(|v| v.as_array()) {
        Some(arr) => arr
            .iter()
            .filter_map(|v| v.as_str().map(str::to_owned))
            .collect(),
        None => return error_frame("invalid_request", "batch-ack requires 'ids' array"),
    };

    match storage.batch_ack(&queue, conn_id, &ids).await {
        Ok(result) => {
            for _ in &result.acked {
                rate_tracker.record_ack(&queue);
            }
            json!({"type": "ok", "acked": result.acked.len()})
        }
        Err(e) => error_frame("batch_ack_failed", &e.to_string()),
    }
}

/// WS-0014, WS-0021: batch negative-acknowledge.
async fn handle_batch_nack(
    frame: &Value,
    storage: &Arc<dyn StorageApi>,
    rate_tracker: &Arc<QueueRateTracker>,
    conn_id: &str,
) -> Value {
    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q.to_owned(),
        None => return error_frame("invalid_request", "batch-nack requires 'queue' field"),
    };
    let ids: Vec<String> = match frame.get("ids").and_then(|v| v.as_array()) {
        Some(arr) => arr
            .iter()
            .filter_map(|v| v.as_str().map(str::to_owned))
            .collect(),
        None => return error_frame("invalid_request", "batch-nack requires 'ids' array"),
    };

    match storage.batch_nack(&queue, conn_id, &ids).await {
        Ok(result) => {
            let nacked = result.unlocked.len() + result.dead_lettered.len() + result.dropped.len();
            for _ in 0..nacked {
                rate_tracker.record_nack(&queue);
            }
            json!({
                "type": "ok",
                "unlocked": result.unlocked.len(),
                "dropped": result.dead_lettered.len() + result.dropped.len(),
            })
        }
        Err(e) => error_frame("batch_nack_failed", &e.to_string()),
    }
}

/// WS-0020: renew (extend) a message lock.
async fn handle_renew(frame: &Value, storage: &Arc<dyn StorageApi>, conn_id: &str) -> Value {
    let (queue, id) = match required_queue_id(frame) {
        Ok(v) => v,
        Err(e) => return e,
    };
    let lock_timeout_secs = frame
        .get("lock_timeout_secs")
        .and_then(|v| v.as_u64())
        .unwrap_or(30);
    match storage.renew(&queue, &id, conn_id, lock_timeout_secs).await {
        Ok(true) => json!({"type": "ok"}),
        Ok(false) => error_frame("not_locked", "message not locked by this consumer"),
        Err(e) => error_frame("renew_failed", &e.to_string()),
    }
}

/// WS-0023: replenish credits for a subscription.
fn handle_credit(frame: &Value, subs: &SubscriptionRegistry) -> Value {
    let queue = match frame.get("queue").and_then(|v| v.as_str()) {
        Some(q) => q,
        None => return error_frame("invalid_request", "credit requires 'queue' field"),
    };
    let credits = match frame.get("credits").and_then(|v| v.as_u64()) {
        Some(c) if c > 0 => c,
        _ => {
            return error_frame(
                "invalid_request",
                "credit requires positive 'credits' field",
            )
        }
    };
    if subs.add_credits(queue, credits) {
        json!({"type": "ok"})
    } else {
        error_frame(
            "not_subscribed",
            &format!("not subscribed to queue '{}'", queue),
        )
    }
}

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

fn error_frame(code: &str, message: &str) -> Value {
    json!({"type": "error", "code": code, "message": message})
}

fn attach_req_id(mut frame: Value, req_id: Option<String>) -> Value {
    if let Some(rid) = req_id {
        if let Some(obj) = frame.as_object_mut() {
            obj.insert("req_id".to_owned(), Value::String(rid));
        }
    }
    frame
}

// ---------------------------------------------------------------------------
// Log subscription handlers
// ---------------------------------------------------------------------------

async fn handle_subscribe_logs(
    log_buffer: &Option<crate::log_buffer::LogBuffer>,
    log_task: &mut Option<tokio::task::JoinHandle<()>>,
    log_tx: &mpsc::Sender<crate::log_buffer::LogEntry>,
) -> Value {
    let buf = match log_buffer {
        Some(b) => b.clone(),
        None => return error_frame("not_available", "log streaming not configured"),
    };

    // If already subscribed, abort old task first.
    if let Some(task) = log_task.take() {
        task.abort();
    }

    let tx = log_tx.clone();

    // History replay and live-event forwarding both run inside the spawned
    // task so that this function returns immediately. Doing the replay inline
    // would deadlock the main select! loop whenever history has more than
    // `log_tx` capacity worth of entries (it fills the mpsc, and the main
    // loop can't drain `log_rx` while it's blocked here). That was
    // reproducible by navigating away from the UI Logs tab and back after
    // enough history had accumulated.
    //
    // Subscribe to the broadcast BEFORE snapshotting history so we can't miss
    // events that arrive during replay. This can produce a small number of
    // duplicates (an entry present in both the history snapshot and the
    // broadcast backlog); duplicates are preferred over missed events, and
    // the UI tolerates them.
    let mut rx = buf.subscribe();
    *log_task = Some(tokio::spawn(async move {
        let history = buf.history_sync();
        for entry in history {
            if tx.send(entry).await.is_err() {
                return; // connection gone
            }
        }
        loop {
            match rx.recv().await {
                Ok(entry) => {
                    if tx.send(entry).await.is_err() {
                        break; // channel closed, connection gone
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
                    // Skip lost entries and continue
                    continue;
                }
                Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
            }
        }
    }));

    json!({"type": "ok"})
}

fn handle_unsubscribe_logs(log_task: &mut Option<tokio::task::JoinHandle<()>>) -> Value {
    if let Some(task) = log_task.take() {
        task.abort();
    }
    json!({"type": "ok"})
}

fn required_queue_id(frame: &Value) -> Result<(String, String), Value> {
    let queue = frame
        .get("queue")
        .and_then(|v| v.as_str())
        .ok_or_else(|| error_frame("invalid_request", "requires 'queue' field"))?
        .to_owned();
    let id = frame
        .get("id")
        .and_then(|v| v.as_str())
        .ok_or_else(|| error_frame("invalid_request", "requires 'id' field"))?
        .to_owned();
    Ok((queue, id))
}