multiplexed-connection 0.4.0

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use cs_trace::{Tracer, child};
use cs_utils::futures::{GenericCodec, wait};
use futures::{StreamExt, SinkExt};
use jsonrpc_core::IoHandler;
use jsonrpc_core_client::transports;
use tokio::io::{AsyncRead, AsyncWrite, DuplexStream};
use tokio::sync::{Mutex, oneshot};
use tokio::sync::mpsc::{Sender, Receiver};
use tokio_util::codec::Framed;
use anyhow::{Result, bail};

use crate::multiplexed_connection::disconnected::RpcMessage;
use connection_utils::{Channel, Connected, TOnRemoteChannelReader};

use crate::utils::AsyncDuplexRpcServer;

use super::rpc::{RpcChannelsServiceClient, RpcChannelRecord, RpcChannelsService, setup_channel, rpc_channel};
use rpc_channel::{RpcChannel, RemoteRpcChannelEvent};


async fn handle_on_rpc_channel_events(
    trace: Box<dyn Tracer>,
    mut on_remote_channel_source: Receiver<RemoteRpcChannelEvent>,
    on_rpc_channel_sink: Sender<Box<dyn Channel>>,
    rpc_client: Arc<RpcChannelsServiceClient>,
) {
    loop {
        if let Ok(channel) = on_remote_channel_source.try_recv() {
            let rpc_channel = RpcChannel::new(
                channel,
                Arc::clone(&rpc_client),
            );

            trace.trace(
                &format!("got rpc channel: \"{:?}\"", &rpc_channel),
            );
            
            let result = on_rpc_channel_sink
                .send(Box::new(rpc_channel)).await;

            // TODO: handle error
            if let Err(error) = result {
                panic!("Cannot send rpc_channel notification: {:?}.", &error.to_string());
            }
            
            wait(10).await;
        }

        wait(20).await;
    }
}

async fn pipe_rpc_messages<TAsyncDuplex: AsyncRead + AsyncWrite + Send + 'static>(
    stream: Pin<Box<TAsyncDuplex>>,
    client_stream: DuplexStream,
    server_stream: DuplexStream,
    mut disconnect_signal_receiver: oneshot::Receiver<()>,
) -> Result<()> {
    let framed_stream = Framed::new(stream, GenericCodec::<RpcMessage>::new());
    let (sink, mut source) = framed_stream.split();

    let client_framed_stream = Framed::new(client_stream, GenericCodec::<String>::new());
    let (mut client_sink, mut client_source) = client_framed_stream.split();

    let server_framed_stream = Framed::new(server_stream, GenericCodec::<String>::new());
    let (mut server_sink, mut server_source) = server_framed_stream.split();

    let sink = Arc::new(Mutex::new(sink));

    let sink1 = Arc::clone(&sink);
    let sink2 = Arc::clone(&sink);

    tokio::try_join!(
        tokio::spawn(async move {
            loop {
                if let Ok(_) = disconnect_signal_receiver.try_recv() {
                    break;
                }

                wait(10).await;
            }

            panic!("Disconnected.");
        }),
        tokio::spawn(async move {
            while let Some(message) = source.next().await {
                let message = message
                    .expect("Failed to get the next RPC message.");
        
                match message {
                    RpcMessage::Client(msg) => {
                        client_sink
                            .send(msg).await
                            .expect("Failed to send client message.");
                    },
                    RpcMessage::Server(msg) => {
                        server_sink
                            .send(msg).await
                            .expect("Failed to send server message.");
                    },
                };
            }
        }),
        // send from the client to the server
        tokio::spawn(async move {
            while let Some(message) = client_source.next().await {
                let message = message
                    .expect("Failed to get the next client message.");
        
                sink1
                    .lock().await
                    .send(RpcMessage::Server(message)).await
                    .expect("Cannot send client message.");
            }
        }),
        // send from the server to the client
        tokio::spawn(async move {
            while let Some(message) = server_source.next().await {
                let message = message
                    .expect("Failed to get the next server message.");
        
                sink2
                    .lock().await
                    .send(RpcMessage::Client(message)).await
                    .expect("Cannot send server message.");
            }
        }),
    )?;

    return Ok(());
}

pub struct MultiplexedConnectionConnected {
    trace: Box<dyn Tracer>,
    channels: Arc<Mutex<HashMap<u16, RpcChannelRecord>>>,
    rpc_client: Arc<RpcChannelsServiceClient>,
    on_remote_data_channel_source: Option<Receiver<Box<dyn Channel>>>,
    disconnect_signal_sender: oneshot::Sender<()>,
}

impl MultiplexedConnectionConnected {
    pub fn new<TAsyncDuplex: AsyncRead + AsyncWrite + Send + 'static>(
        trace: &Box<dyn Tracer>,
        channels: Arc<Mutex<HashMap<u16, RpcChannelRecord>>>,
        rpc_client: Arc<RpcChannelsServiceClient>,
        on_remote_data_channel_source: Option<Receiver<Box<dyn Channel>>>,
        on_remote_channel_source: Receiver<RemoteRpcChannelEvent>,
        on_rpc_channel_sink: Sender<Box<dyn Channel>>,
        stream: Pin<Box<TAsyncDuplex>>,
        client_duplex_sink: DuplexStream,
        server_duplex_sink: DuplexStream,
    ) -> Self {
        let (
            disconnect_signal_sender,
            disconnect_signal_receiver,
        ) = oneshot::channel::<()>();

        // TODO: propagate errors?
        tokio::spawn(
            handle_on_rpc_channel_events(
                child!(trace, "on-channel"),
                on_remote_channel_source,
                on_rpc_channel_sink,
                Arc::clone(&rpc_client),
            ),
        );

        // TODO: propagate errors?
        tokio::spawn(
            pipe_rpc_messages(
                stream,
                client_duplex_sink,
                server_duplex_sink,
                disconnect_signal_receiver,
            ),
        );

        let trace = child!(trace, "rpc-con-connected");
        return MultiplexedConnectionConnected {
            trace,
            channels,
            rpc_client,
            on_remote_data_channel_source,
            disconnect_signal_sender,
        };
    }
}

impl MultiplexedConnectionConnected {
    pub async fn connect_client<TAsyncDuplex>(
        trace: &Box<dyn Tracer>,
        client_stream: TAsyncDuplex,
    ) -> Result<RpcChannelsServiceClient>
        where
            TAsyncDuplex: AsyncRead + AsyncWrite + Send + 'static,
    {
        let framed_remote_stream = Framed::new(client_stream, GenericCodec::<String>::new());
        let (sink, mut source) = framed_remote_stream.split();

        let stream = async_stream::stream! {
            while let Some(maybe_message) = source.next().await {
                let message = maybe_message
                    .expect("Failed to get next message.");

                yield message;
            }
        };

        let (rpc_client_future, sender) = transports::duplex(
            Box::pin(sink), 
            Box::pin(stream),
        );
        let rpc_client = RpcChannelsServiceClient::from(sender);

        // TODO: should propagate errors correctly?
        tokio::spawn(async move {
            rpc_client_future.await.unwrap();
        });

        trace.trace("rpc client created");

        return Ok(rpc_client);
    }

    pub async fn connect_server<TAsyncDuplex>(
        trace: &Box<dyn Tracer>,
        channels: Arc<Mutex<HashMap<u16, RpcChannelRecord>>>,
        server_stream: TAsyncDuplex,
        on_data_channel_sink: Sender<RemoteRpcChannelEvent>,
    ) -> Result<AsyncDuplexRpcServer>
        where
            TAsyncDuplex: AsyncRead + AsyncWrite + Send + 'static
    {
        let mut server_io = IoHandler::new();
        server_io.extend_with(
            RpcChannelsService::new(
                trace,
                channels,
                on_data_channel_sink,
            ).as_delegate(),
        );

        let server = AsyncDuplexRpcServer::new(server_io)
            .build(server_stream);

        return Ok(server);
    }
}

#[async_trait]
impl Connected for MultiplexedConnectionConnected {
    fn on_remote_channel(&mut self) -> Result<TOnRemoteChannelReader> {
        let source = match self.on_remote_data_channel_source.take() {
            Some(source) => source,
            None => bail!("No on_remote_channel source found, already taken?"),
        };

        return Ok(source);
    }

    // TODO: test
    fn off_remote_channel(
        &mut self,
        on_remote_data_channel_source: TOnRemoteChannelReader,
    ) -> Result<()> {
        if self.on_remote_data_channel_source.is_some() {
            bail!("Remote channel source already set.");
        }

        self.on_remote_data_channel_source = Some(on_remote_data_channel_source);

        return Ok(());
    }

    /// Create local data channel. In this case we reuse the `RpcChannelsService`(local_server)
    /// logic to create channel, minus firing the `on_remote_channel` event.
    async fn channel(&mut self, label: String) -> Result<Box<dyn Channel>> {
        self.trace.debug(
            &format!("creating channel"),
        );

        // create channel on the remote side
        let channel_id = match self.rpc_client
            .create_channel(label.clone()).await
        {
            Ok(id) => id,
            Err(error) => {
                bail!("Remote channel creation failed: {:?}", &error);
            },
        };

        // create channel record in the local channels map
        let channel = setup_channel(
            &self.trace,
            Some(channel_id),
            label,
            Arc::clone(&self.channels),
        ).await?;

        self.trace.trace(
            &format!("channel set up, creating rpc_channel"),
        );

        let rpc_channel = RpcChannel::new(
            channel,
            Arc::clone(&self.rpc_client),
        );

        return Ok(Box::new(rpc_channel));
    }

    /// Disconnect.
    /// TODO: must send disconnect RPC notification to the other side?
    async fn disconnect(mut self) -> Result<()> {
        if let Err(_) = &self.disconnect_signal_sender.send(()) {
            bail!("Failed to send disconnect signal.");
        }

        return Ok(());
    }
}

#[cfg(test)]
mod test {
    mod disconnect {
        // use rstest::rstest;
        // use tokio::io::duplex;
        // use tokio::try_join;

        // use cs_utils::futures::wait;
        // use cs_utils::{random_str_rg};
        // use cs_utils::futures::test::test_async_stream;

        // use crate::MultiplexedConnection;

        // #[rstest]
        // #[case::size_8_32(8, 32)]
        // #[case::size_128_512(128, 512)]
        // #[case::size_2048_4096(2048, 4096)]
        // #[case::size_4096_8192(4096, 8192)]
        // #[case::size_8192_16384(8192, 16384)]
        // #[tokio::test]
        // async fn sends_data_from_remote_channel(
        //     #[case] str_min_size: usize,
        //     #[case] str_max_size: usize,
        // ) {
        //     let (duplex1, duplex2) = duplex(4096);

        //     let channel_label = format!("channel-label-{}", cs_utils::random_str(4));
        //     let channel_label1 = channel_label.clone();
        //     let channel_label2 = channel_label.clone();

        //     let (channel1, channel2) = try_join!(
        //         tokio::spawn(async move {
        //             let mut connection1 = MultiplexedConnection::new(duplex1)
        //                 .connect().await
        //                 .expect("Error while listening");

        //             cs_utils::futures::wait(50).await;
                    
        //             let channel = connection1
        //                 .channel(channel_label1.clone()).await
        //                 .unwrap();

        //             assert_eq!(
        //                 channel.label(),
        //                 &channel_label1,
        //                 "Channel labels must match.",
        //             );

        //             channel
        //         }),
        //         tokio::spawn(async move {
        //             let mut connection2 = MultiplexedConnection::new(duplex2)
        //                 .listen().await
        //                 .expect("Error while listening.");

        //             let mut on_remote_channel = connection2.on_remote_channel().unwrap();

        //             let channel = {
        //                 loop {
        //                     if let Ok(channel) = on_remote_channel.try_recv() {
        //                         break channel;
        //                     }

        //                     wait(50).await;
        //                 }
        //             };

        //             assert_eq!(
        //                 channel.label(),
        //                 &channel_label2,
        //                 "Channel labels must match.",
        //             );

        //             channel
        //         }),
        //     ).unwrap();

        //     let test_data = vec![
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //         random_str_rg(str_min_size..=str_max_size),
        //     ].join("");

        //     test_async_stream(
        //         channel1,
        //         channel2,
        //         test_data,
        //     ).await;
        // }
    }
}