rtx 0.1.0

use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

use pin_project_lite::pin_project;

use crate::net::ToSocketAddrs;

pin_project! {
    /// A TCP stream between a local and a remote socket.
    #[derive(Debug)]
    pub struct TcpStream {
        #[pin]
        inner: FromRuntime! {
            "std" => std::net::TcpStream,
            // https://docs.rs/async-std/latest/src/async_std/net/tcp/stream.rs.html#50
            "async-std" => async_io::Async<std::net::TcpStream>,
            "tokio" => tokio::net::TcpStream,
        },
    }
}

impl TcpStream {
    /// Opens a TCP connection to a remote host in non-blocking mode.
    #[cfg(feature = "async")]
    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
    pub async fn connect<A: ToSocketAddrs>(address: A) -> io::Result<Self> {
        let addrs = address.to_socket_addrs().await?;
        let mut last_error = None;

        for addr in addrs {
            let r = pick_async_runtime_async!(addr, {
                "async-std" => |addr| async_io::Async::<std::net::TcpStream>::connect(addr),
                "tokio" => |addr| tokio::net::TcpStream::connect(addr),
            })
            .resolve()
            .await
            .into_result();

            match r {
                Ok(stream) => {
                    return Ok(Self { inner: stream });
                }

                Err(error) => {
                    last_error = Some(error);
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "could not resolve to any addresses",
            )
        }))
    }

    /// Opens a TCP connection to a remote host in blocking mode.
    #[cfg(any(feature = "blocking", doc))]
    #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
    pub fn connect_blocking<A: ToSocketAddrs>(address: A) -> io::Result<Self> {
        let addrs = address.to_socket_addrs_blocking()?;
        let mut last_error = None;

        for addr in addrs {
            let r = pick_blocking_runtime!(addr, {
                "std" => |addr| std::net::TcpStream::connect(addr),
            });

            match r {
                Ok(stream) => {
                    return Ok(Self { inner: stream });
                }

                Err(error) => {
                    last_error = Some(error);
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "could not resolve to any addresses",
            )
        }))
    }

    /// Gets the value of the `TCP_NODELAY` option on this socket.
    pub fn nodelay(&self) -> io::Result<bool> {
        apply_runtime!(&self.inner, {
            "std" => |stream| stream.nodelay(),
            "async-std" => |stream| stream.get_ref().nodelay(),
            "tokio" => |stream| stream.nodelay(),
        })
    }

    /// Sets the value of the `TCP_NODELAY` option on this socket.
    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
        apply_runtime!(&self.inner, {
            "std" => |stream| stream.set_nodelay(nodelay),
            "async-std" => |stream| stream.get_ref().set_nodelay(nodelay),
            "tokio" => |stream| stream.set_nodelay(nodelay),
        })
    }
}

/// # Note: Blocking Behavior Follows Construction
/// If `TcpStream` was created with [`Self::connect_blocking()`],
/// this impl will use blocking I/O.
///
/// Otherwise, this impl will return `std::io::ErrorKind::WouldBlock` instead of blocking, and
/// then you can use the [`Ready`] trait to wait for the stream to be ready to read again.
/// This behavior is intended as it makes it useful with existing APIs based on blocking I/O
/// but which can handle `WouldBlock`, such as [`rustls`](https://crates.io/crates/rustls).
impl io::Read for TcpStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        apply_mut_runtime!(&mut self.inner, buf, {
            "std" => |stream, buf| stream.read(buf),
            "async-std" => |stream, buf| stream.get_mut().read(buf),
            "tokio" => |stream, buf| stream.try_read(buf),
        })
    }
}

/// # Note: Blocking Behavior Follows Construction
/// If `TcpStream` was created with [`Self::connect_blocking()`],
/// this impl will use blocking I/O.
///
/// Otherwise, this impl will return `std::io::ErrorKind::WouldBlock` instead of blocking, and
/// then you can use the [`Ready`] trait to wait for the stream to be ready to write again.
/// This behavior is intended as it makes it useful with existing APIs based on blocking I/O
/// but which can handle `WouldBlock`, such as [`rustls`](https://crates.io/crates/rustls).
impl io::Write for TcpStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        apply_mut_runtime!(&mut self.inner, {
            "std" => |stream| stream.write(buf),
            "async-std" => |stream| stream.get_mut().write(buf),
            "tokio" => |stream| stream.try_write(buf),
        })
    }

    fn flush(&mut self) -> io::Result<()> {
        // NOTE: flush on TCP streams is a nop
        Ok(())
    }
}

impl crate::io::Ready for TcpStream {
    fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        apply_runtime!(&self.inner, cx, {
            "std" => |_, _| Poll::Ready(Ok(())),

            "async-std" => |stream, cx| {
                Pin::new(stream).poll_readable(cx)
            },

            "tokio" => |stream, cx| {
                Pin::new(stream).poll_read_ready(cx)
            },
        })
    }

    fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        apply_runtime!(&self.inner, cx, {
            "std" => |_, _| Poll::Ready(Ok(())),

            "async-std" => |stream, cx| {
                Pin::new(stream).poll_writable(cx)
            },

            "tokio" => |stream, cx| {
                Pin::new(stream).poll_write_ready(cx)
            },
        })
    }
}

#[cfg(feature = "async")]
impl crate::io::AsyncRead for TcpStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        apply_mut_runtime_async!(&mut self.project().inner, (cx, buf), {
            "async-std" => |stream, (cx, buf)| {
                Pin::new(stream).poll_read(cx, buf)
            },

            "tokio" => |stream, (cx, buf)| {
                use tokio::io::{AsyncRead, ReadBuf};

                let mut rbuf = ReadBuf::new(buf);

                futures_core::ready!(Pin::new(stream).poll_read(cx, &mut rbuf)?);

                Poll::Ready(Ok(rbuf.filled().len()))
            },
        })
    }
}

#[cfg(feature = "async")]
impl crate::io::AsyncWrite for TcpStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        apply_mut_runtime_async!(&mut self.project().inner, (cx, buf), {
            "async-std" => |stream, (cx, buf)| {
                Pin::new(stream).poll_write(cx, buf)
            },

            "tokio" => |stream, (cx, buf)| {
                use tokio::io::AsyncWrite;

                Pin::new(stream).poll_write(cx, buf)
            },
        })
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        // NOTE: flush on TCP streams is a nop
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        apply_mut_runtime_async!(&mut self.project().inner, cx, {
            "async-std" => |stream, cx| {
                Pin::new(stream).poll_close(cx)
            },

            "tokio" => |stream, cx| {
                use tokio::io::AsyncWrite;

                Pin::new(stream).poll_shutdown(cx)
            },
        })
    }
}