tubes 0.6.4

use crate::client_id::ClientId;
use crate::{data::MessageDataInternal, prelude::*};

mod transport;

use pipenet::{NonBlockStream, Packs};
use std::{
    collections::{HashMap, HashSet},
    io::{Error, ErrorKind},
    net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, TcpStream},
    sync::{
        Arc, Mutex,
        mpsc::{Receiver, Sender, channel},
    },
    thread::JoinHandle,
    time::{Duration, Instant},
};

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

/// A session: intended as a network session across a set of nodes.
///
/// Networking is established through TCP sockets.
///
/// Internally using [pipenet](https://docs.rs/pipenet).
///
/// One of this nodes will act as a server and relay messages to the rest of
/// the nodes, the clients. The server needs to be started first before
/// clients can connect.
///
/// The instance of this session will maintain a unique and random (v4) [`ClientId`]
/// on creation and it will not change until the session instance is dropped.
/// Even when host promotion happens the uuids are maintained. This will keep
/// the concept of the session alive and each node can rely on node uuids to
/// remain stable within the same session.
///
/// Each instance handles the session for the point of view of each node,
/// including the handles to background threads, channels and internal buffers.
/// The instance for the client keeps one background thread to pipe the I/O
/// into its channels, while the server also has an additional thread that will
/// loop on TCP connection `accept`. The [Config] can specify a timeout for the
/// server accepting new clients, to not block further clients connecting after
/// that.
///
/// Dropping this instance closes the related connection(s): when client it
/// will disconnect, when server it will also disconnect all the other clients.
/// It is however also possible to manually disconnect a session with
/// [`Session::stop`]. The session can then be started again.
///
/// Messages can be sent to all nodes using [`Session::broadcast`], or to a
/// specific one with [`Session::send_to`] using the destination's [`ClientId`].
///
/// Receiving of messages is done through [`Session::read`]. Each of those calls
/// are non-blocking and will return [Some] in case there is a message
/// available. The message is wrapped in [`MessageData`] and can represent also
/// a few more useful extra messages provided by this implementation, such as
/// the ones that allow for [`ClientId`] identification of nodes when joining or
/// leaving. See example.
///
/// Sessions can migrate their host using [`Session::promote_to_host`]. This has
/// to be called on the node that is currently the server. This can take some
/// time and is not immediate. More over, the final stage of migration is
/// triggered only when the session is interacted by the user, during one of
/// the read/write operations, such as [`Session::broadcast`],
/// [`Session::send_to`], or [`Session::read`]. Keep polling the session after
/// requesting a promotion to ensure the full stage is completed. During this
/// phase messages to be sent are held back during send/broadcast and sent only
/// after the reconnection has happened. While reading instead, the promotion
/// will trigger only when there are no more messages in the queue to be
/// consumed by [`Session::read`].
///
/// The type `MS` is specifying the maximum size of a packet, that if exceeded,
/// the connection will close to the client. Default is 0, which means no
/// limits.
///
/// server and client example:
/// ```rust
/// use tubes::prelude::*;
/// use std::thread::sleep;
/// use std::time::Duration;
/// use std::string::FromUtf8Error;
///
/// #[derive(Clone, Debug, PartialEq)]
/// struct Msg(String);
///
/// impl From<String> for Msg {
///     fn from(value: String) -> Self {
///         Self(value)
///     }
/// }
///
/// impl TryFrom<&[u8]> for Msg {
///     type Error = FromUtf8Error;
///
///     fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> {
///         Ok(Msg(String::from_utf8(value.to_vec())?))
///     }
/// }
///
/// impl TryFrom<Msg> for Vec<u8> {
///     type Error = ();
///
///     fn try_from(value: Msg) -> std::result::Result<Self, Self::Error> {
///         Ok(value.0.into())
///     }
/// }
///
/// let mut s = Session::<0>::new_server(":5000".into());
/// s.start().unwrap();
///
/// let mut c = Session::<0>::new_client("127.0.0.1:5000".into());
/// c.start().unwrap();
///
/// assert!(s.is_connected());
/// assert!(c.is_connected());
/// println!("Connected.");
/// sleep(Duration::from_millis(100));
/// // Server internally knows the list of all its clients, by uuid.
/// for uuid in s.clients() {
///     println!("Client is: {}", uuid);
/// }
///
/// s.broadcast("hello".to_string().into());
/// sleep(Duration::from_millis(100));
/// if let MessageData::Broadcast{from, data} = c.read().unwrap().unwrap() {
///     println!("Message from {from}: {data:?}");
/// }
///
/// c.stop();
/// s.stop();
/// ```
pub struct Session<const MS: usize = 0> {
    kind: SessionKind,
    config: Config,
    accept_routine: Option<JoinHandle<()>>,
    io_routine: Option<JoinHandle<()>>,
    uuid: ClientId,
    server_uuid: Option<ClientId>,
    clients: Arc<Mutex<HashMap<ClientId, NonBlockStream<MS>>>>,
    tx_writer: Option<Sender<MessageDataInternal>>,
    rx_reader: Option<Arc<Mutex<Receiver<MessageData>>>>,
    reconnect_to: Arc<Mutex<Option<ReconnectTo>>>,
}

impl<const MS: usize> Session<MS> {
    /// Creates a new server from the configuration
    #[must_use]
    pub fn new_server(config: Config) -> Self {
        Self {
            kind: SessionKind::Server,
            config,
            accept_routine: None,
            io_routine: None,
            uuid: ClientId::new(),
            server_uuid: None,
            clients: Arc::default(),
            tx_writer: None,
            rx_reader: None,
            reconnect_to: Mutex::new(None).into(),
        }
    }

    /// Creates a new server from the configuration
    #[must_use]
    pub fn new_client(config: Config) -> Self {
        Self {
            kind: SessionKind::Client,
            config,
            accept_routine: None,
            io_routine: None,
            uuid: ClientId::new(),
            server_uuid: None,
            clients: Arc::default(),
            tx_writer: None,
            rx_reader: None,
            reconnect_to: Mutex::new(None).into(),
        }
    }

    /// Starts the session.
    /// If server, binds to the port,
    /// If client, connects to the address.
    ///
    /// Starting a started session is a no operation.
    ///
    /// Spawns the necessary background threads.
    ///
    /// # Errors
    ///
    /// Error thrown if can't connect.
    pub fn start(&mut self) -> Result<()> {
        match self.kind {
            SessionKind::Server => self.start_server(),
            SessionKind::Client => self.start_client(),
        }
    }

    /// Stops the connection.
    /// If server, closes also all the clients,
    /// If client, stops the current connection.
    ///
    /// Stopping a stopped session is a no operation.
    ///
    /// Every thread is terminated and handles removed.
    pub fn stop(&mut self) {
        self.accept_routine = None;
        self.io_routine = None;
        self.rx_reader = None;
        self.tx_writer = None;
    }

    /// Returns if the current session is a server, client otherwise.
    ///
    /// A session started as server could become a client later and vice versa,
    /// when the host promotion happens and one of the clients transitions to
    /// become the server of the session.
    #[must_use]
    pub fn is_server(&self) -> bool {
        self.kind == SessionKind::Server
    }

    /// This returns true when the background thread is active on an open
    /// stream, or open server if it is a server.
    #[must_use]
    pub fn is_connected(&self) -> bool {
        match self.kind {
            SessionKind::Server => {
                if let Some(h) = &self.accept_routine
                    && !h.is_finished()
                    && let Some(h) = &self.io_routine
                    && !h.is_finished()
                {
                    return true;
                }
                false
            }
            SessionKind::Client => {
                if let Some(h) = &self.io_routine
                    && !h.is_finished()
                {
                    return true;
                }
                false
            }
        }
    }

    /// Returns which uuid is the server
    ///
    /// This may also return [None] if no connection is established.
    ///
    /// If server, returns self uuid.
    #[must_use]
    pub fn server_uuid(&self) -> Option<ClientId> {
        if self.is_server() {
            Some(self.uuid)
        } else {
            self.server_uuid
        }
    }

    /// Returns the uuid of this endpoint.
    /// For clients it's the client uuid, for servers is the server uuid.
    #[must_use]
    pub fn uuid(&self) -> ClientId {
        self.uuid
    }

    /// Returns the clients currently connected to this server, if this is a
    /// server, otherwise it's an empty slice.
    ///
    /// Clients are identified by uuid.
    #[must_use]
    pub fn clients(&self) -> HashSet<ClientId> {
        let Ok(lock) = self.clients.lock() else {
            return HashSet::new();
        };
        lock.keys().copied().collect::<HashSet<ClientId>>()
    }

    /// Reads for one message from the internal channel.
    /// This operation does not block.
    ///
    /// If one message is returned, chances are there are more available to
    /// be read, so call this method in a loop as long as it returns Some.
    ///
    /// This is a no operation when the node is disconnected.
    ///
    /// # Errors
    ///
    /// This can error in case of broken pipe. Session needs to be discarded
    /// after that, or reconnected.
    pub fn read(&mut self) -> Result<Option<MessageData>> {
        if !self.is_connected() {
            return Err(Error::new(ErrorKind::NotConnected, "not connected").into());
        }
        let Some(c) = self.rx_reader.as_mut() else {
            return Ok(None);
        };
        let Ok(c) = c.lock() else {
            return Ok(None);
        };
        let msg = c.try_recv().ok();
        drop(c);
        if msg.is_none() {
            // Only attempt reconnection when the queue has been emptied,
            // otherwise the reader may miss some messages.
            self.check_reconnect_to()?;
        }
        Ok(msg)
    }

    /// Sends the message only to that specific node, by uuid.
    /// Server automatically redirects this to the destination.
    ///
    /// This is a no operation when the node is disconnected.
    ///
    /// # Errors
    ///
    /// This can error in case of broken pipe. Session needs to be discarded
    /// after that, or reconnected.
    pub fn send_to(&mut self, uuid: ClientId, m: Vec<u8>) -> Result<()> {
        if !self.is_connected() {
            return Err(Error::new(ErrorKind::NotConnected, "not connected").into());
        }
        // Only attempt reconnection beore sending messages.
        self.check_reconnect_to()?;
        let Some(c) = self.tx_writer.as_mut() else {
            return Ok(());
        };
        let _ = c.send(MessageDataInternal::Send(self.uuid, uuid, m));
        Ok(())
    }

    /// Sends the message to all clients except self.
    /// If client, only send to server as broadcast,
    /// If server, it will consume and repeat also to all the other clients.
    ///
    /// This is a no operation when the node is disconnected.
    ///
    /// # Errors
    ///
    /// This can error in case of broken pipe. Session needs to be discarded
    /// after that, or reconnected.
    pub fn broadcast(&mut self, m: Vec<u8>) -> Result<()> {
        if !self.is_connected() {
            return Err(Error::new(ErrorKind::NotConnected, "not connected").into());
        }
        // Only attempt reconnection beore sending messages.
        self.check_reconnect_to()?;
        let Some(c) = self.tx_writer.as_mut() else {
            return Ok(());
        };
        let _ = c.send(MessageDataInternal::Broadcast(self.uuid, m));
        Ok(())
    }

    /// Promote the given uuid to become the new server.
    /// This sends the messages to begin promotion.
    ///
    /// The first stage happens in the background where all the nodes are sent
    /// notification of the transition and they become ready for it.
    ///
    /// During the normal routines such as [`Session::send_to`],
    /// [`Session::broadcast`] or [`Session::read`] each node will trigger the
    /// second stage where the socket is actually recreated and the background
    /// threads spawned as new.
    ///
    /// This must happen in this way because only the main accessor of the
    /// [Session] can trigger a thread & socket recreation. This is similar to
    /// calling [`Session::stop`], changing the internal addressing, and then
    /// calling [`Session::start`].
    ///
    /// The original uuids of each node are maintained after the promotion, as
    /// the [Session] instances stay the same, they just reconnect to a new
    /// topology.
    ///
    /// This is a no operation when the node is disconnected or it is not a
    /// server node (only servers can promote clients).
    ///
    /// * `uuid` - the uuid of the client that will become the server.
    /// * `port` - which port to use, pass None to use the same of the server.
    ///
    pub fn promote_to_host(&mut self, uuid: ClientId, port: Option<u16>) {
        if !self.is_server() {
            return;
        }
        let Some(c) = self.tx_writer.as_mut() else {
            return;
        };
        let Ok(map) = self.clients.lock() else {
            return;
        };
        let Some(client) = map.get(&uuid) else {
            return;
        };
        // Find what ip and port the client has, from the perspective of the
        // server: this allows to determine a bind address that is more valid,
        // since it was at least proven to work point-to-point as a client,
        // because the socket stream was connected to it.
        // (This will not work if the client is behind NAT).
        let addr = client.remote_addr().ip();
        let port = port.unwrap_or(self.config.port);
        let msg = MessageDataInternal::PromoteToHost(uuid, addr, port);
        let _ = c.send(msg);
    }

    /// Returns the total amount of bytes read.
    /// If this was a server session, it is the sum of all the bytes read from
    /// all clients.
    #[must_use]
    pub fn total_read(&self) -> usize {
        let mut t = 0;
        let Ok(map) = self.clients.lock() else {
            return 0;
        };
        for (_, c) in map.iter() {
            t += c.total_read();
        }
        t
    }

    /// Returns the total amount of bytes sent.
    /// If this was a server session, it is the sum of all the bytes sent to
    /// all clients.
    #[must_use]
    pub fn total_sent(&self) -> usize {
        let mut t = 0;
        let Ok(map) = self.clients.lock() else {
            return 0;
        };
        for (_, c) in map.iter() {
            t += c.total_sent();
        }
        t
    }

    fn start_server(&mut self) -> Result<()> {
        if self.is_connected() {
            return Ok(());
        }

        let addr = self
            .config
            .address
            .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED));

        let addr = SocketAddr::from((addr, self.config.port));
        let server_uuid = self.uuid;
        let listener = TcpListener::bind(addr)?;
        let accept_timeout = self.config.accept_timeout;
        let client_list = self.clients.clone();
        let config = self.config.clone();
        self.accept_routine = Some(std::thread::spawn(move || {
            transport::accept_loop(server_uuid, &config, &client_list, &listener, accept_timeout);
        }));

        let server_uuid = self.uuid;
        let client_list = self.clients.clone();
        let reconnect_to = self.reconnect_to.clone();
        let (tx_reader, rx_reader) = channel();
        let (tx_writer, rx_writer) = channel();
        self.tx_writer = Some(tx_writer);
        self.rx_reader = Some(Mutex::new(rx_reader).into());
        self.io_routine = Some(std::thread::spawn(move || {
            transport::server_loop::<MS>(
                server_uuid,
                &client_list,
                &reconnect_to,
                &rx_writer,
                &tx_reader,
            );
        }));

        Ok(())
    }

    fn start_client(&mut self) -> Result<()> {
        if self.is_connected() {
            return Ok(());
        }

        let Some(addr) = self.config.address else {
            return Ok(());
        };

        let addr = SocketAddr::from((addr, self.config.port));
        // Retry a few times, it could be in between a host promotion so it can
        // take sometime to the new server to start.
        let socket = connect_with_retry_and_wait(addr)?;
        socket.set_nonblocking(true)?;
        let mut socket = to_pipenet::<MS>(socket, &self.config);
        // Important: must send the uuid of this client right away or the
        // server will wait for it until accept_timeout, which may slow down
        // the other accepts queued on the line.
        socket.write((MessageDataInternal::ClientJoined(self.uuid)).try_into()?)?;
        let Some(server_uuid) =
            wait_for_server_uuid_message(self.config.accept_timeout, &mut socket)?
        else {
            return Err("Could not connect to server: did not receive server uuid.".into());
        };
        self.server_uuid = Some(server_uuid);

        if let Ok(mut map) = self.clients.lock() {
            map.clear();
            map.insert(self.uuid, socket.clone());
        }

        let reconnect_to = self.reconnect_to.clone();
        let (tx_reader, rx_reader) = channel();
        let (tx_writer, rx_writer) = channel();
        self.tx_writer = Some(tx_writer);
        self.rx_reader = Some(Mutex::new(rx_reader).into());
        self.io_routine = Some(std::thread::spawn(move || {
            transport::client_loop::<MS>(socket, &reconnect_to, &rx_writer, &tx_reader);
        }));

        Ok(())
    }

    // This method must reconnect synchronously because it may be called just
    // before queueing a new message and that requires the channels to be
    // alive. The connection may come later, but the channels are synchronously
    // recreated on reconnection.
    fn check_reconnect_to(&mut self) -> Result<()> {
        let Ok(mut reconnect_to) = self.reconnect_to.lock() else {
            return Ok(());
        };
        let Some(ref to) = *reconnect_to else {
            return Ok(());
        };
        let server = to.become_server;
        let address = to.address;
        let port = to.port;
        *reconnect_to = None;
        drop(reconnect_to);

        self.stop();

        self.config.address = Some(address);
        self.config.port = port;
        self.kind = if server {
            SessionKind::Server
        } else {
            SessionKind::Client
        };

        self.start()
    }
}

impl<const MS: usize> Drop for Session<MS> {
    fn drop(&mut self) {
        if let Some(c) = self.tx_writer.as_ref() {
            let _ = c.send(MessageDataInternal::ClientLeft(self.uuid));
        }
        self.stop();
    }
}

#[derive(Default, PartialEq)]
enum SessionKind {
    #[default]
    Server,
    Client,
}

fn connect_with_retry_and_wait(addr: SocketAddr) -> Result<TcpStream> {
    let mut ct = 0;
    loop {
        match TcpStream::connect(addr) {
            Ok(stream) => return Ok(stream),
            Err(e) => {
                if ct > 10 {
                    return Err(e.into());
                }
                std::thread::sleep(Duration::from_millis(100));
                ct += 1;
            }
        }
    }
}

pub(crate) struct ReconnectTo {
    // Use this to determine when the newserver(as old client) is gone from the
    // old server. The uuids stay the same in host promotion.
    // And when the uuid is the same as self (Session::uuid()) and self is a
    // client, then this is a client that is asked to become a server for this
    // address binding & port.
    pub(crate) become_server: bool,
    pub(crate) address: IpAddr,
    pub(crate) port: u16,
}

pub(crate) fn to_pipenet<const MS: usize>(
    stream: TcpStream,
    config: &Config,
) -> NonBlockStream<MS> {
    #[allow(unused_mut)]
    let mut packs = Packs::default();
    #[cfg(feature = "compression")]
    if config.compress {
        packs = packs.compress();
    }
    #[cfg(feature = "encryption")]
    if let Some(key) = config.key.as_ref() {
        packs = packs.encrypt(key);
    }
    NonBlockStream::<MS>::from_version_packs(config.versions, packs, stream)
}

fn wait_for_server_uuid_message<const MS: usize>(
    timeout: Duration,
    client: &mut NonBlockStream<MS>,
) -> Result<Option<ClientId>> {
    let now = Instant::now();
    loop {
        let Some(msg) = client.read()? else {
            continue;
        };
        let msg = MessageDataInternal::try_from(msg.as_slice())?;
        if let MessageDataInternal::ServerUuid(uuid) = msg {
            return Ok(Some(uuid));
        }
        if now.elapsed() > timeout {
            return Ok(None);
        }
    }
}