rustwave 0.1.0

use std::{future::Future, pin::Pin, sync::Arc, time::Duration};

use dashmap::DashMap;
use futures_util::{SinkExt, StreamExt};
use serde::Serialize;
use tokio::{
    sync::{Notify, mpsc::Sender},
    time::sleep,
};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tokio_util::sync::CancellationToken;

use crate::{
    error::Result,
    utils::send_data::{Payload, Request, Response},
};

type EventType = Arc<
    DashMap<String, Arc<dyn Fn(Payload) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>,
>;

/// A WebSocket client that connects to a Server (or any
/// compatible WebSocket server) and exchanges events with it.
///
/// # Example
///
/// ```no_run
/// use rustwave::{Client, Result};
///
/// #[tokio::main]
/// async fn main() -> Result<()> {
///     let client = Client::connect("ws://127.0.0.1:8080").await?;
///
///     client.on("message", |data| async move {
///         println!("Server say: {}", data);
///     });
///
///     client.emit("message", "hello").await?;
///     client.wait().await;
///
///     Ok(())
/// }
/// ```
#[derive(Clone)]
pub struct Client {
    // Channel used to queue outgoing messages, which the internal write
    // task forwards to the actual socket (single writer for the connection).
    sender: Sender<Message>,

    // Signals shutdown across the client's background tasks (read, write,
    // heartbeat). Cancelling one stops the others too.
    token: CancellationToken,

    // Registered event handlers.
    // key = event name, value = the handler to run when it's received.
    events: EventType,
}

impl Client {
    /// Connects to a WebSocket server at the given `url`.
    ///
    /// On success, this establishes the connection and spawns background
    /// tasks to handle reading incoming events, sending outgoing events, and
    /// keeping the connection alive with periodic heartbeat pings. Event
    /// handlers can then be registered with [`Client::on`] and messages sent
    /// with [`Client::emit`].
    ///
    /// # Errors
    ///
    /// Returns an error if the connection to `url` cannot be established.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Client, Result};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let client = Client::connect("ws://127.0.0.1:8080").await?;
    ///
    ///     client.on("message", |data| async move {
    ///         println!("Server say: {}", data);
    ///     });
    ///
    ///     Ok(())
    /// }
    /// ```
    pub async fn connect(url: &str) -> Result<Self> {
        let (stream, _) = connect_async(url).await?;
        let (mut write, mut read) = stream.split();

        let events: EventType = Arc::new(DashMap::new());

        events.insert(
            "error".to_string(),
            Arc::new(|payload| {
                Box::pin(async move {
                    eprintln!("error: {}", payload.raw());
                })
            }),
        );
        let read_event = events.clone();

        let (tx, mut rx) = tokio::sync::mpsc::channel::<Message>(100);
        let read_tx = tx.clone();
        let heartbeat = tx.clone();

        let token = CancellationToken::new();
        let read_token = token.clone();
        let wait_token = token.clone();
        let write_token = token.clone();
        let heartbeat_token = token.clone();

        let notify = Arc::new(Notify::new());

        let heartbeat_notify = notify.clone();
        let read_notify = notify.clone();

        tokio::spawn(async move {
            let heartbeat = heartbeat;
            loop {
                tokio::select! {
                    _ = sleep(Duration::from_secs(30)) => {}
                    _ = heartbeat_token.cancelled() => break
                }

                let _ = heartbeat.send(Message::Ping(Vec::new().into())).await;

                tokio::select! {
                    res = tokio::time::timeout(Duration::from_secs(10), heartbeat_notify.notified()) => {
                        if res.is_err() {
                            let _ = heartbeat.send(Message::Close(None)).await;
                            heartbeat_token.cancel();
                            break
                        }
                    }

                    _ = heartbeat_token.cancelled() => break
                }
            }
        });

        tokio::spawn(async move {
            let read_event = read_event.clone();
            loop {
                tokio::select! {
                    mass = read.next() => {
                        match mass {
                            Some(Ok(Message::Text(text))) => {

                                let parsed = match serde_json::from_str::<Request>(&text) {
                                    Ok(v) => v,
                                    Err(_) => continue
                                };

                                if let Some(handler) = read_event.get(&parsed.event) {
                                    let payload = Payload::new(parsed.data);
                                    handler(payload).await;
                                }
                            }
                            Some(Ok(Message::Close(_))) => {
                                read_token.cancel();

                                break
                            }
                            Some(Ok(Message::Ping(data))) => {
                                let _ = read_tx.send(Message::Pong(data)).await;
                            }
                            Some(Ok(Message::Pong(_))) => {
                                read_notify.notify_one();
                            },

                            Some(Ok(_)) => (),
                            None | Some(Err(_)) => {
                                read_token.cancel();
                                break
                            }
                        }
                    }

                    _ = read_token.cancelled() => break
                }
            }

            if let Some(handler) = read_event.get("disconnect") {
                handler(Payload::new(serde_json::Value::Null)).await;
            }
        });

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    mass = rx.recv() => {
                        match mass {
                            Some(m) => {
                                let is_close = matches!(m, Message::Close(_));
                                match write.send(m).await {
                                    Ok(_v) => {
                                        if is_close {
                                            write_token.cancel();
                                            break;
                                        }
                                    }
                                    Err(_) => {
                                        write_token.cancel();
                                        break
                                    }
                                }
                            }
                            None => {
                                write_token.cancel();
                                break
                            }
                        }
                    }

                    _ = write_token.cancelled() => break
                }
            }
        });

        Ok(Self {
            sender: tx,
            token: wait_token,
            events,
        })
    }

    /// Registers an event handler on this client.
    ///
    /// Registering a handler for an event that already has one replaces the
    /// previous handler.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Client, Result};
    ///
    /// # #[tokio::main]
    /// # async fn example() -> Result<()>{
    ///     let client = Client::connect("ws://127.0.0.1:8080").await?;
    ///
    ///     client.on("message", |data| async move {
    ///         println!("Server say: {}", data);
    ///     });
    ///     Ok(())
    /// # }
    /// ```
    pub fn on<F, Fut>(&self, event: &str, callback: F)
    where
        F: Fn(Payload) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.events
            .insert(event.to_string(), Arc::new(move |v| Box::pin(callback(v))));
    }

    /// Sends an event to the server.
    ///
    /// `data` is serialized to JSON and sent over the connection.
    ///
    /// # Errors
    ///
    /// Returns an error if `data` fails to serialize, or if the message
    /// could not be sent (e.g. the connection is closed).
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Client, Result};
    ///
    /// # async fn example() -> Result<()> {
    ///     let client = Client::connect("ws://127.0.0.1:8080").await?;
    ///
    ///     client.emit("message", "hello").await?;
    ///
    ///     Ok(())
    /// # }
    /// ```
    pub async fn emit<T: Serialize>(&self, event: &str, data: T) -> Result<()> {
        let json = serde_json::to_string(&Response {
            event: event.to_string(),
            data,
        })?;

        match self.sender.send(Message::Text(json.into())).await {
            Ok(_) => Ok(()),
            Err(_) => Err(crate::error::RustWaveError::Internal(
                "Message sending error".to_string(),
            )),
        }
    }

    /// Closes the WebSocket connection.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Client, Result};
    ///
    /// # async fn example() -> Result<()> {
    ///     let client = Client::connect("ws://127.0.0.1:8080").await?;
    ///
    ///     client.close().await;
    ///     Ok(())
    /// # }
    /// ```
    pub async fn close(&self) {
        let _ = self.sender.send(Message::Close(None)).await;
    }

    /// Waits until the client's connection has been closed.
    ///
    /// This resolves once the underlying connection has shut down — whether
    /// due to calling [`Client::close`], the server closing the connection,
    /// a heartbeat timeout, or a connection error. Useful for keeping a
    /// program alive while the client runs in the background.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Client, Result};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let client = Client::connect("ws://127.0.0.1:8080").await?;
    ///
    ///     client.on("message", |data| async move {
    ///         println!("Server say: {}", data);
    ///     });
    ///
    ///   
    ///     client.wait().await;
    ///     Ok(())
    /// }
    /// ```
    pub async fn wait(&self) {
        self.token.cancelled().await;
        println!("Disconnected");
    }
}