rustwave 0.1.0

use dashmap::DashMap;
use futures_util::{SinkExt, StreamExt};
use serde_json::Value;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::time::sleep;
use tokio_tungstenite::tungstenite::handshake::server::Request;

use tokio_tungstenite::accept_hdr_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::sync::CancellationToken;

use crate::error::Result;
use crate::peer::peer_conext::PeerContext;
use crate::server::namespace::Namespace;
use crate::utils::send_data::{Payload, Request as WsRequest};

/// A WebSocket server. that listens for incoming connections and routes
/// them to registered [`Namespace`]s based on the connection's URL path.
///
/// Every server has an implicit **root namespace** (`"/"`).
/// Calling `Server.on` registers events on the root namespace.
///
/// # Example
/// ```no_run
/// use rustwave::{Server, Result};
///
/// #[tokio::main]
/// async fn main() -> Result<()> {
///    let mut server = Server::bind("127.0.0.1:8080");
///    
///    // Events on the root namespace("/")
///    server.on("message", |data, ctx| async move {
///        println!("Client {} say: {}", ctx.get_socket_id(), data);
///    });
///
///     // Events on the custom namespace
///     let chat = server.namespace("/chat");
///     chat.on("message", |data, ctx| async move {
///         println!("Client {} say: {}", ctx.get_socket_id(), data);
///     });
///
///     // run the server
///     server.run().await?;
///     Ok(())
/// }
///
/// ```
pub struct Server {
    // Address the server will bind to, e.g. "127.0.0.1:8080"
    addr: String,

    // Registry of all namespaces.
    // key = path ("/", "/chat"), value = the Namespace for that path.
    namespace: Arc<DashMap<String, Namespace>>,

    // The default/root namespace ("/")
    // Clients that connect without specifying a path land here
    // Server.on() registers events on this root namespace directly
    root: Namespace,
}

impl Server {
    /// Create a new server bound to the specified address.
    /// This does not start listening for incoming connections.
    ///
    /// The `addr` parameter spaecifies the address server will bind to.
    /// for example: `127.0.0.1:8080` or `0.0.0.0:8080`
    /// bind return type is `Server`
    ///
    /// # Example
    /// ```no_run
    /// use rustwave::{ Server, Result };
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     Ok(())
    /// }
    /// ```
    pub fn bind(addr: &str) -> Self {
        let namespace: Arc<DashMap<String, Namespace>> = Arc::new(DashMap::new());
        let root = Namespace::new();

        namespace.insert("/".to_string(), root.clone());

        Self {
            addr: addr.to_string(),
            namespace,
            root,
        }
    }

    // get all namespaces
    // Return a handle to the server's namespace registry.
    // This can be used to access namespaces registry. its managed by the server
    pub(crate) fn get_namespace(&self) -> Arc<DashMap<String, Namespace>> {
        self.namespace.clone()
    }

    /// Returns the [`Namespace`] for the given path.
    ///
    /// if the namespace does not already exist, it is created and returned.
    ///
    /// this is normalized to always start with `/` so `chat` and `/chat` refer the same namespace
    ///
    /// Calling this meethod multiple times with the same path will always return the same namespace.
    /// Events on that namespace are shared.
    ///
    /// Clients connect to a namespace by including its path in the WebSocket URL.
    /// for example: `ws://localhost:8080/chat`
    ///
    /// # Example
    /// ```no_run
    /// use rustwave::{ Server, Result };
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///     chat.on("message", |data, ctx| async move {
    ///         println!("Client {} say: {}", ctx.get_socket_id(), data);
    ///     });
    ///     Ok(())
    /// }
    /// ```
    ///
    /// # Arguments
    ///
    /// * `path` - The path of the namespace to retrieve.   
    ///
    pub fn namespace(&self, path: &str) -> Namespace {
        let name = if path.starts_with('/') {
            path.to_string()
        } else {
            format!("/{}", path)
        };
        self.namespace
            .entry(name.to_string())
            .or_insert_with(Namespace::new)
            .value()
            .clone()
    }

    /// Register an event handler on the server's root namespace (`"/"`).
    ///
    /// This is convenient method equivalent to calling `namespace("/").on(event, handler)`
    ///
    /// Event handlers registered with this method are invoked only for clients
    /// that connect to the root namespace (`"/"`) — i.e. clients that connect
    /// without specifying any path (e.g. `ws://localhost:8080/`).
    ///
    /// # Arguments
    ///
    /// * `event` - The name of the event to register the handler for.
    /// * `handler` - The handler function that will be invoked when the event is triggered.
    ///
    /// # Example
    /// ```no_run
    /// use rustwave::{ Server, Result };
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     server.on("message", |data, ctx| async move {
    ///         println!("Client {} say: {}", ctx.get_socket_id(), data);
    ///     });
    ///     Ok(())
    /// }
    /// ```
    pub fn on<F, Fut>(&self, event: &str, handler: F)
    where
        F: Fn(Payload, PeerContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.root.on(event, handler);
    }

    /// Starts the server.
    ///
    /// Calling this method server begins listening for incoming WebSocket connections.
    ///
    /// Register all event handlers and namespaces before calling this method.
    /// This method runs until the server is shut down or an unrecoverable
    /// error occurs.
    ///
    /// # Errors
    ///
    /// Returns an error if the server fails to bind to the configured address.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{ Server, Result };
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     server.run().await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn run(&self) -> Result<()> {
        let listener = TcpListener::bind(&self.addr).await?;

        loop {
            let (stream, _) = listener.accept().await?;

            let namespace = self.get_namespace();

            tokio::spawn(async move {
                let mut path = String::new();

                let ws = match accept_hdr_async(stream, |req: &Request, res| {
                    path = req.uri().path().to_string();
                    Ok(res)
                })
                .await
                {
                    Ok(ws) => ws,
                    Err(_) => return,
                };

                let (mut write, mut read) = ws.split();
                let (tx, mut rx) = tokio::sync::mpsc::channel::<Message>(100);

                let ns = match namespace.get(&path) {
                    Some(ns) => ns.clone(),
                    None => {
                        let _ = write
                            .send(Message::Text("This path is not found".into()))
                            .await;
                        let _ = write.send(Message::Close(None)).await;
                        return;
                    }
                };

                let context = ns.create_peer(tx.clone());
                let client = context.get_peer_state();

                ns.handler("connection", Value::Null, context.clone());

                let socket_id = client.get_socket_id();
                ns.insert_client(socket_id, tx.clone());

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

                let heartbeat = tx.clone();
                let notify = Arc::new(Notify::new());

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

                let heartbeat_task = 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
                        }
                    }
                });

                let tx_sender = tx.clone();
                let client_rooms = client.get_rooms();

                let read_task = tokio::spawn(async move {
                    let emitter = context;
                    let sender = tx_sender;
                    let rooms = client_rooms;

                    loop {
                        tokio::select! {
                            mass = read.next() => {
                                match mass {
                                    Some(m) => {
                                        match m {
                                            Ok(message) => {
                                                match message {
                                                    Message::Text(text) => {
                                                        let parsed = match serde_json::from_str::<WsRequest>(&text) {
                                                            Ok(v) => v,
                                                            Err(_) => continue
                                                        };
                                                        ns.handler(&parsed.event, parsed.data, emitter.clone());
                                                    }

                                                    Message::Ping(data) => {
                                                        match sender.send(Message::Pong(data)).await {
                                                            Ok(_) => {}
                                                            Err(_) => break
                                                        }
                                                    }

                                                    Message::Close(frame) => {
                                                        match sender.send(Message::Close(frame)).await {
                                                            Ok(_) => break,
                                                            Err(_) => break
                                                        }
                                                    }

                                                    Message::Pong(_) => {
                                                        read_notify.notify_one();
                                                    }

                                                    _ => {}
                                                }
                                            }

                                            Err(_) => break
                                        }
                                    }
                                    None => break
                                }
                            }

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

                    ns.remove_client(&socket_id);

                    for room in rooms.iter() {
                        ns.remove_room(&room, &socket_id);
                    }

                    read_token.cancel();

                    ns.handler("disconnect", Value::Null, emitter.clone());
                });

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

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

                let _ = tokio::join!(read_task, write_task, heartbeat_task);
            });
        }
    }
}