ws_reconnect_client/

client.rs

use backon::{BackoffBuilder, ExponentialBuilder};
use futures_util::{SinkExt, StreamExt};
use serde::{de::DeserializeOwned, Serialize};
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;

use crate::{
    connect_with_retry, MessageStream, PingManager, Result, WebSocketError, WsConnectionConfig,
    WsReader, WsWriter,
};

/// High-level WebSocket client with automatic reconnection, ping/pong handling, and JSON support
#[derive(Clone)]
pub struct WebSocketClient<T>
where
    T: DeserializeOwned,
{
    config: WsConnectionConfig,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> WebSocketClient<T>
where
    T: DeserializeOwned,
{
    /// Create a new WebSocket client with the given configuration
    pub fn new(config: WsConnectionConfig) -> Self {
        Self {
            config,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Connect to the WebSocket server without sending any subscription
    ///
    /// Returns a writer and reader for manual message handling.
    pub async fn connect(&self) -> Result<(WsWriter, WsReader)> {
        connect_with_retry(&self.config).await
    }

    /// Connect to the WebSocket server and return a MessageStream
    ///
    /// The MessageStream automatically handles PING/PONG and yields parsed messages.
    pub async fn connect_stream(&self) -> Result<MessageStream<T>> {
        use std::sync::Arc;
        use tokio::sync::Mutex;

        let (writer, reader) = self.connect().await?;
        let shared_writer = Arc::new(Mutex::new(Some(writer)));
        Ok(MessageStream::new(
            reader,
            shared_writer,
            self.config.ping_interval_secs,
        ))
    }

    /// Connect and optionally send an initial subscription message
    ///
    /// Deprecated: Use `connect()` + `send_subscription()` or `connect_stream()` instead.
    pub async fn connect_and_subscribe<S: Serialize>(
        &self,
        subscription: Option<&S>,
    ) -> Result<(WsWriter, WsReader)> {
        let (mut writer, reader) = self.connect().await?;

        if let Some(sub) = subscription {
            send_subscription(&mut writer, sub).await?;
        }

        Ok((writer, reader))
    }

    /// Start listening to messages with automatic reconnection and ping/pong handling
    ///
    /// Uses exponential backoff for reconnection attempts (same config as initial connection).
    ///
    /// # Arguments
    /// * `subscription` - Optional subscription message to send on each connection
    /// * `handler` - Callback function that processes each received message
    ///
    /// # Example
    /// ```no_run
    /// # use ws_reconnect_client::{WebSocketClient, WsConnectionConfig};
    /// # use serde::{Deserialize, Serialize};
    /// # #[derive(Debug, Deserialize, Serialize)]
    /// # struct MyMessage { value: String }
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let config = WsConnectionConfig::new("wss://example.com");
    /// # let subscription = MyMessage { value: "sub".into() };
    /// let client = WebSocketClient::<MyMessage>::new(config);
    /// client.listen(Some(&subscription), |msg| {
    ///     println!("Received: {:?}", msg);
    ///     Ok(())
    /// }).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn listen<S, F>(&self, subscription: Option<S>, mut handler: F) -> Result<()>
    where
        S: Serialize + Clone,
        F: FnMut(T) -> Result<()>,
    {
        if !self.config.auto_reconnect {
            // No auto-reconnect, just listen once
            return self.listen_once(subscription.as_ref(), &mut handler).await;
        }

        // Auto-reconnect enabled: use exponential backoff for reconnections
        let backoff = ExponentialBuilder::default()
            .with_min_delay(Duration::from_millis(self.config.initial_backoff_ms))
            .with_max_delay(Duration::from_millis(self.config.max_backoff_ms))
            .with_max_times(self.config.max_retries);

        let mut backoff_iter = backoff.build();
        let mut attempt = 0;

        loop {
            match self.listen_once(subscription.as_ref(), &mut handler).await {
                Ok(_) => {
                    // Connection closed normally, reset backoff and reconnect
                    backoff_iter = backoff.build();
                    attempt = 0;
                }
                Err(e) => {
                    // Connection error
                    attempt += 1;

                    if attempt >= self.config.max_retries {
                        // Max retries reached
                        return Err(e);
                    }

                    // Get next backoff delay from backon
                    if let Some(delay) = backoff_iter.next() {
                        tokio::time::sleep(delay).await;
                    } else {
                        // Backoff exhausted (shouldn't happen due to max_retries check)
                        return Err(e);
                    }
                }
            }
        }
    }

    /// Listen once (single connection lifecycle)
    async fn listen_once<S, F>(&self, subscription: Option<&S>, handler: &mut F) -> Result<()>
    where
        S: Serialize,
        F: FnMut(T) -> Result<()>,
    {
        let (mut writer, mut reader) = self.connect_and_subscribe(subscription).await?;

        let mut ping_manager = if self.config.ping_interval_secs > 0 {
            Some(PingManager::new(self.config.ping_interval_secs))
        } else {
            None
        };

        loop {
            tokio::select! {
                // Handle incoming messages
                msg = reader.next() => {
                    match msg {
                        Some(Ok(Message::Text(text))) => {
                            self.handle_text_message(&text, handler)?;
                        }

                        Some(Ok(Message::Ping(ping))) => {
                            writer.send(Message::Pong(ping))
                                .await
                                .map_err(|_| WebSocketError::SendError)?;
                        }

                        Some(Ok(Message::Pong(_))) => {
                            // Pong received, connection is alive
                        }

                        Some(Ok(Message::Close(_))) => {
                            return Err(WebSocketError::ConnectionClosed);
                        }

                        Some(Err(e)) => {
                            return Err(WebSocketError::Tungstenite(e));
                        }

                        None => {
                            return Err(WebSocketError::ConnectionClosed);
                        }

                        _ => {}
                    }
                }

                // Send periodic pings if enabled
                _ = async {
                    if let Some(ref mut pm) = ping_manager {
                        pm.wait_for_next_ping().await;
                    } else {
                        // Never completes if pings are disabled
                        std::future::pending::<()>().await;
                    }
                } => {
                    writer.send(Message::Ping(vec![].into()))
                        .await
                        .map_err(|_| WebSocketError::SendError)?;
                }
            }
        }
    }

    /// Handle text message - try to parse as single message or array of messages
    fn handle_text_message<F>(&self, text: &str, handler: &mut F) -> Result<()>
    where
        F: FnMut(T) -> Result<()>,
    {
        // Skip empty messages (keepalive frames from some servers)
        if text.trim().is_empty() {
            return Ok(());
        }

        // Try to parse as array first
        if text.trim_start().starts_with('[') {
            match serde_json::from_str::<Vec<T>>(text) {
                Ok(messages) => {
                    for msg in messages {
                        handler(msg)?;
                    }
                    return Ok(());
                }
                Err(e) => {
                    // Array parsing failed, log the error and try single message
                    eprintln!("⚠️  WebSocket: Failed to parse as array: {}", e);
                }
            }
        }

        // Try to parse as single message
        match serde_json::from_str::<T>(text) {
            Ok(msg) => {
                handler(msg)?;
                Ok(())
            }
            Err(e) => {
                // Could be a subscription confirmation or other non-T message
                // Log raw message for debugging
                eprintln!("⚠️  WebSocket: Failed to parse message: {}", e);
                eprintln!("📨 Raw message: {}", text);
                // Don't fail, just skip it
                Err(WebSocketError::SerializationError(e))
            }
        }
    }
}

/// Send a subscription message over an existing WebSocket connection
///
/// This is a helper function for sending JSON-serialized subscription messages.
///
/// # Example
/// ```no_run
/// # use ws_reconnect_client::{WsWriter, send_subscription};
/// # #[derive(serde::Serialize)]
/// # struct Subscription { channel: String }
/// # async fn example(writer: &mut WsWriter) -> Result<(), Box<dyn std::error::Error>> {
/// let sub = Subscription { channel: "market".to_string() };
/// send_subscription(writer, &sub).await?;
/// # Ok(())
/// # }
/// ```
pub async fn send_subscription<S: Serialize>(writer: &mut WsWriter, subscription: &S) -> Result<()> {
    let sub_json = serde_json::to_string(subscription)?;
    writer
        .send(Message::Text(sub_json.into()))
        .await
        .map_err(|_| WebSocketError::SendError)
}

/// Builder pattern for WebSocketClient
pub struct WebSocketClientBuilder {
    config: WsConnectionConfig,
}

impl WebSocketClientBuilder {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            config: WsConnectionConfig::new(url),
        }
    }

    pub fn with_config(config: WsConnectionConfig) -> Self {
        Self { config }
    }

    pub fn ping_interval(mut self, seconds: u64) -> Self {
        self.config = self.config.with_ping_interval(seconds);
        self
    }

    pub fn auto_reconnect(mut self, enabled: bool) -> Self {
        self.config = self.config.with_auto_reconnect(enabled);
        self
    }

    pub fn max_retries(mut self, retries: usize) -> Self {
        self.config = self.config.with_retries(retries);
        self
    }

    pub fn backoff(mut self, initial_ms: u64, max_ms: u64) -> Self {
        self.config = self.config.with_backoff(initial_ms, max_ms);
        self
    }

    pub fn build<T: DeserializeOwned>(self) -> WebSocketClient<T> {
        WebSocketClient::new(self.config)
    }
}