use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::{TcpStream, UnixStream};
use tokio_listener::Connection;

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

pub enum AbstractStream {
    Tcp(TcpStream),
    Uds(UnixStream),
    Listener(Connection),
}

impl From<TcpStream> for AbstractStream {
    fn from(stream: TcpStream) -> Self {
        Self::Tcp(stream)
    }
}

impl From<UnixStream> for AbstractStream {
    fn from(stream: UnixStream) -> Self {
        Self::Uds(stream)
    }
}

impl From<Connection> for AbstractStream {
    fn from(stream: Connection) -> Self {
        Self::Listener(stream)
    }
}

impl Unpin for AbstractStream {}

impl AsyncRead for AbstractStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<()>> {
        match &mut *self {
            Self::Tcp(stream) => pin!(stream).poll_read(cx, buf),
            Self::Uds(stream) => pin!(stream).poll_read(cx, buf),
            Self::Listener(stream) => pin!(stream).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for AbstractStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize>> {
        match &mut *self {
            Self::Tcp(stream) => pin!(stream).poll_write(cx, buf),
            Self::Uds(stream) => pin!(stream).poll_write(cx, buf),
            Self::Listener(stream) => pin!(stream).poll_write(cx, buf),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        match &mut *self {
            Self::Tcp(stream) => pin!(stream).poll_flush(cx),
            Self::Uds(stream) => pin!(stream).poll_flush(cx),
            Self::Listener(stream) => pin!(stream).poll_flush(cx),
        }
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
        match &mut *self {
            Self::Tcp(stream) => pin!(stream).poll_shutdown(cx),
            Self::Uds(stream) => pin!(stream).poll_shutdown(cx),
            Self::Listener(stream) => pin!(stream).poll_shutdown(cx),
        }
    }
}