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

use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

pub enum Socket {
    Tcp(tokio::net::TcpStream),
    #[cfg(unix)]
    Unix(tokio::net::UnixStream),
}

impl From<tokio::net::TcpStream> for Socket {
    fn from(socket: tokio::net::TcpStream) -> Self {
        Self::Tcp(socket)
    }
}

#[cfg(unix)]
impl From<tokio::net::UnixStream> for Socket {
    fn from(socket: tokio::net::UnixStream) -> Self {
        Self::Unix(socket)
    }
}

impl AsyncRead for Socket {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buffer: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        match self.get_mut() {
            Self::Tcp(socket) => Pin::new(socket).poll_read(cx, buffer),
            #[cfg(unix)]
            Self::Unix(socket) => Pin::new(socket).poll_read(cx, buffer),
        }
    }
}

impl AsyncWrite for Socket {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buffer: &[u8],
    ) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            Self::Tcp(socket) => Pin::new(socket).poll_write(cx, buffer),
            #[cfg(unix)]
            Self::Unix(socket) => Pin::new(socket).poll_write(cx, buffer),
        }
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buffers: &[IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            Self::Tcp(socket) => Pin::new(socket).poll_write_vectored(cx, buffers),
            #[cfg(unix)]
            Self::Unix(socket) => Pin::new(socket).poll_write_vectored(cx, buffers),
        }
    }

    fn is_write_vectored(&self) -> bool {
        match self {
            Self::Tcp(socket) => socket.is_write_vectored(),
            #[cfg(unix)]
            Self::Unix(socket) => socket.is_write_vectored(),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match self.get_mut() {
            Self::Tcp(socket) => Pin::new(socket).poll_flush(cx),
            #[cfg(unix)]
            Self::Unix(socket) => Pin::new(socket).poll_flush(cx),
        }
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        match self.get_mut() {
            Self::Tcp(socket) => Pin::new(socket).poll_shutdown(cx),
            #[cfg(unix)]
            Self::Unix(socket) => Pin::new(socket).poll_shutdown(cx),
        }
    }
}