twfss/src/network/client.rs

use std::time::Duration;

use async_tungstenite::{tokio::connect_async, tungstenite::Message, WebSocketStream};
use core::pin::Pin;
use futures::future::{self, Either, Future};
use futures::task::{Context, Poll};
use tokio::stream::Stream;
use tokio::stream::StreamExt;
use tokio::sync::{mpsc, oneshot};
use tokio::time::{delay_for, Delay};
use url::Url;

use super::packet::IncomingPacket;

pub fn create<StreamType, ErrorType, PacketType>(
    stream: StreamType,
) -> (
    RPCInterface<StreamType, PacketType>,
    ServerEventStream<PacketType>,
)
where
    StreamType: Stream<Item = Result<Vec<u8>, ErrorType>> + Unpin + Send + 'static,
{
    let (cmd_send, cmd_receive) = mpsc::channel(1);
    tokio::spawn(async move {
        // TODO: Remove type annotations.
        client_connection_task::<StreamType, ErrorType, PacketType>(stream, cmd_receive).await;
    });
    // TODO
    panic!("Not yet implemented.");
}

enum ClientTaskCommand {
    Close,
    SendPackets,
}

struct OutgoingCallMap {
    // TODO
}

async fn client_connection_task<StreamType, ErrorType, PacketType>(
    mut stream: StreamType,
    mut cmd_receive: mpsc::Receiver<ClientTaskCommand>,
) where
    StreamType: Stream<Item = Result<Vec<u8>, ErrorType>> + Unpin,
{
    let mut connection_closed = false;
    loop {
        match future::select(Box::pin(stream.next()), Box::pin(cmd_receive.recv())).await {
            Either::Left((packet, _cmd_future)) => {
                match packet {
                    Some(data) => {
                        // TODO
                    }
                    None => {
                        // Stream closed. We still need to return errors for any
                        // packet sent by the client, so do not exit the loop
                        // immediately but rather return an error for any
                        // incoming packet.
                        connection_closed = true;
                        break;
                    }
                }
            }
            Either::Right((cmd, _stream_future)) => {
                match cmd {
                    Some(ClientTaskCommand::Close) => {
                        // The RPC interface for this connection was dropped, so
                        // we just exit the loop and close the connection.

                        // We need to return errors for all pending packets
                        // first, though.
                        // TODO
                        break;
                    }
                    Some(ClientTaskCommand::SendPackets) => {
                        // TODO
                    }
                    None => {
                        // This should never happen.
                        break;
                    }
                }
            }
        };
    }

    if connection_closed {
        // The stream was closed, but the command channel was not. We still need
        // to return errors for any packet sent by the client, so do not exit
        // the loop immediately but rather return an error for any incoming
        // packet.
        client_connection_closed(cmd_receive).await;
    } else {
        // The command channel was closed, but the stream was not. Gracefully
        // close the stream.
        // TODO
    }
}

async fn client_connection_closed(mut cmd_receive: mpsc::Receiver<ClientTaskCommand>) {
    loop {
        match cmd_receive.recv().await {
            Some(ClientTaskCommand::Close) => {
                // The connection was already closed.
                break;
            }
            Some(ClientTaskCommand::SendPackets) => {
                // The connection was closed, so return an error for any
                // packet sent.
                // TODO
            }
            None => {
                // This should never happen.
                break;
            }
        }
    }
}

pub struct RPCInterface<StreamType, PacketType> {
    stream: StreamType,
    _packet: Option<PacketType>, // TODO: Remove.
}

impl<StreamType, PacketType> RPCInterface<StreamType, PacketType> {
    fn call(_call: &PacketType) -> RPC<StreamType, PacketType> {
        // TODO
        panic!("Not yet implemented.");
    }
}

pub struct RPC<PacketType, ErrorType> {
    result: Pin<Box<oneshot::Receiver<Result<IncomingPacket<PacketType>, ErrorType>>>>,
    timeout: Option<Pin<Box<Delay>>>,
}
impl<PacketType, ErrorType> RPC<PacketType, ErrorType> {
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(Box::pin(delay_for(timeout)));
        self
    }
}

impl<PacketType, ErrorType> Future for RPC<PacketType, ErrorType> {
    type Output = Result<IncomingPacket<PacketType>, RPCError<ErrorType>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        // Safe, as we will not move self_.
        let self_ = unsafe { self.get_unchecked_mut() };

        // Try receiving from the receiver.
        match Pin::as_mut(&mut self_.result).poll(cx) {
            Poll::Ready(Ok(result)) => return Poll::Ready(result.map_err(|x| RPCError::Stream(x))),
            Poll::Ready(Err(_)) => {
                // The channel was closed.
                return Poll::Ready(Err(RPCError::Closed));
            }
            Poll::Pending => (),
        }

        // Nothing received, check the timeout instead.
        if self_.timeout.is_some() {
            match Pin::as_mut(&mut self_.timeout.as_mut().unwrap()).poll(cx) {
                Poll::Ready(()) => return Poll::Ready(Err(RPCError::Timeout)),
                Poll::Pending => (),
            }
        }

        Poll::Pending
    }
}

// TODO: Do we even need this type?
pub struct ServerEventStream<PacketType> {
    packets: Pin<Box<mpsc::Receiver<IncomingPacket<PacketType>>>>,
}

impl<PacketType> Stream for ServerEventStream<PacketType> {
    type Item = IncomingPacket<PacketType>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        // Safe, as we will not move self_.
        let self_ = unsafe { self.get_unchecked_mut() };

        Pin::as_mut(&mut self_.packets).poll_next(cx)
    }
}

pub enum RPCError<StreamError> {
    Closed,
    Timeout,
    Stream(StreamError)
}

pub async fn low_level_integration_test_client(address: String) {
    use super::websocket::WebSocketConnection;

    let url = Url::parse(&address).unwrap();
    let (mut ws_stream, _) = connect_async(url).await.unwrap();
    let mut conn = WebSocketConnection::new(ws_stream);
    let (rpc_interface, server_events) = create::<_, _, u32>(conn);
    // TODO
    panic!("Not yet implemented.");
}