use std::mem::MaybeUninit;
use std::os::windows::io::{AsRawSocket, AsSocket, BorrowedSocket, RawSocket};
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use std::{io, slice};

use socket2::{SockAddr, SockRef};
use tokio::io::{AsyncRead, AsyncWrite, Interest, ReadBuf};

use crate::net::Domain;
use crate::net::general::{
    DgramSpec, ListenerTy, StreamSpec, StreamTy, UniSocket, UnknownTy,
    impl_bound_socket_local_addr, impl_connected_socket_peer_addr, impl_raw_socket_bind,
    impl_raw_socket_new, impl_tcp_listener_accept, impl_tcp_listener_bind,
    impl_tcp_listener_poll_accept, impl_tcp_socket_connect, impl_tcp_socket_listen,
    impl_tcp_stream_connect, impl_tcp_stream_peek, impl_tcp_stream_poll_peek,
    impl_tcp_stream_poll_read_fn, impl_tcp_stream_poll_write_fn, impl_tcp_stream_read,
    impl_tcp_stream_read_fn, impl_tcp_stream_split, impl_tcp_stream_write,
    impl_tcp_stream_write_fn, impl_udp_socket_connect, impl_udp_socket_peek,
    impl_udp_socket_peek_from, impl_udp_socket_peek_sender, impl_udp_socket_poll_peek,
    impl_udp_socket_poll_peek_from, impl_udp_socket_poll_peek_sender, impl_udp_socket_poll_recv,
    impl_udp_socket_poll_recv_fn, impl_udp_socket_poll_recv_from, impl_udp_socket_poll_send,
    impl_udp_socket_poll_send_fn, impl_udp_socket_poll_send_to, impl_udp_socket_recv,
    impl_udp_socket_recv_fn, impl_udp_socket_recv_from, impl_udp_socket_send,
    impl_udp_socket_send_fn, impl_udp_socket_send_to,
};

/// A TCP socket.
pub type TcpSocket = UniSocket<tokio::net::TcpSocket, StreamSpec, UnknownTy>;

/// An established TCP stream (i.e. a TCP socket connected to a remote peer).
pub type TcpStream = UniSocket<tokio::net::TcpStream, StreamSpec, StreamTy>;

/// An owned read half (RX) of a [`TcpStream`].
pub type OwnedReadHalf = UniSocket<tokio::net::tcp::OwnedReadHalf, StreamSpec, StreamTy>;

/// An owned write half (TX) of a [`TcpStream`].
pub type OwnedWriteHalf = UniSocket<tokio::net::tcp::OwnedWriteHalf, StreamSpec, StreamTy>;

/// A TCP socket listening for connections.
pub type TcpListener = UniSocket<tokio::net::TcpListener, StreamSpec, ListenerTy>;

/// A UDP socket.
pub type UdpSocket = UniSocket<tokio::net::UdpSocket, DgramSpec, UnknownTy>;

impl<I, Spec, Ty> AsSocket for UniSocket<I, Spec, Ty>
where
    I: AsSocket,
{
    #[cfg_attr(coverage_nightly, coverage(off))]
    #[inline]
    /// Returns a [`BorrowedSocket`] to the underlying socket.
    ///
    /// This applys to all socket types, including [`TcpSocket`], [`TcpStream`],
    /// [`OwnedReadHalf`], [`OwnedWriteHalf`], [`TcpListener`], and
    /// [`UdpSocket`].
    fn as_socket(&self) -> BorrowedSocket<'_> {
        self.inner.as_socket()
    }
}

impl<I, Spec, Ty> AsRawSocket for UniSocket<I, Spec, Ty>
where
    I: AsRawSocket,
{
    #[cfg_attr(coverage_nightly, coverage(off))]
    #[inline]
    /// Returns the raw socket handle of the underlying socket.
    ///
    /// This applys to all socket types, including [`TcpSocket`], [`TcpStream`],
    /// [`OwnedReadHalf`], [`OwnedWriteHalf`], [`TcpListener`], and
    /// [`UdpSocket`].
    fn as_raw_socket(&self) -> RawSocket {
        self.inner.as_raw_socket()
    }
}

impl<I, Spec, Ty> UniSocket<I, Spec, Ty>
where
    I: AsSocket,
{
    #[cfg_attr(coverage_nightly, coverage(off))]
    #[inline]
    /// Returns a [`SockRef`] to the underlying socket.
    ///
    /// This applys to all socket types, including [`TcpSocket`], [`TcpStream`],
    /// [`OwnedReadHalf`], [`OwnedWriteHalf`], [`TcpListener`], and
    /// [`UdpSocket`].
    ///
    /// ## Notes
    ///
    /// This is not a free operation, since `socket2` does not provide a way to
    /// get a `SockRef` without performing some needless checks.
    pub fn to_ref(&self) -> SockRef<'_> {
        SockRef::from(self)
    }
}

impl TcpSocket {
    impl_raw_socket_new!(|socket| {
        // (this comment is for formatting)
        Ok(tokio::net::TcpSocket::from_std_stream(socket.into()))
    });

    impl_raw_socket_bind!(|self, addr| {
        let Some(addr) = addr.as_socket() else {
            return Err(io::Error::other("unsupported socket address family"));
        };

        self.inner.bind(addr).map(|()| self)
    });

    impl_bound_socket_local_addr!(|self| {
        // (this comment is for formatting)
        self.inner.local_addr().map(From::from)
    });

    impl_tcp_socket_connect!(|self, addr| {
        let Some(addr) = addr.as_socket() else {
            return Err(io::Error::other("unsupported socket address family"));
        };

        self.inner.connect(addr).await.map(TcpStream::from_inner)
    });

    impl_tcp_socket_listen!(|self, backlog| {
        self.inner
            .listen(backlog.cast_unsigned())
            .map(TcpListener::from_inner)
    });
}

impl TcpStream {
    impl_tcp_stream_connect!(|addr| {
        // (this comment is for formatting)
        TcpSocket::new(addr.domain())?.connect(addr).await
    });

    impl_bound_socket_local_addr!(|self| {
        // (this comment is for formatting)
        self.inner.local_addr().map(From::from)
    });

    impl_connected_socket_peer_addr!(|self| {
        // (this comment is for formatting)
        self.inner.peer_addr().map(From::from)
    });

    impl_tcp_stream_peek!(|self, buf| {
        // `tokio` requires the buffer to be initialized. Oh, performance be damned,
        // let's just zero it out.
        buf.fill(MaybeUninit::new(0));

        #[allow(unsafe_code, reason = "XXX")]
        // SAFETY: We just filled the buffer with zeroes, so it is now
        // initialized.
        let buf = unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len()) };

        self.inner.peek(buf).await
    });

    impl_tcp_stream_poll_peek!(|self, cx, buf| {
        self.inner.poll_peek(cx, &mut ReadBuf::uninit(buf))
    });

    impl_tcp_stream_read!(|self, buf| {
        loop {
            self.inner.readable().await?;

            match self.inner.try_read_buf(&mut ReadBuf::uninit(buf)) {
                read @ Ok(_) => break read,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                ret @ Err(_) => return ret,
            }
        }
    });

    impl_tcp_stream_write!(|self, buf| {
        loop {
            self.inner.writable().await?;

            match self.inner.try_write(buf) {
                write @ Ok(_) => break write,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                ret @ Err(_) => return ret,
            }
        }
    });

    impl_tcp_stream_split!(|self| {
        let (r, w) = self.inner.into_split();

        (OwnedReadHalf::from_inner(r), OwnedWriteHalf::from_inner(w))
    });

    impl_tcp_stream_read_fn!(|self, f| {
        let sockref = SockRef::from(&self.inner);

        self.inner.async_io(Interest::READABLE, move || f(&sockref))
    });

    impl_tcp_stream_poll_read_fn!(|self, cx, f| {
        let sockref = SockRef::from(&self.inner);

        loop {
            ready!(self.inner.poll_read_ready(cx))?;

            match self.inner.try_io(Interest::READABLE, || f(&sockref)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });

    impl_tcp_stream_write_fn!(|self, f| {
        let sockref = SockRef::from(&self.inner);

        self.inner.async_io(Interest::WRITABLE, move || f(&sockref))
    });

    impl_tcp_stream_poll_write_fn!(|self, cx, f| {
        let sockref = SockRef::from(&self.inner);

        loop {
            ready!(self.inner.poll_write_ready(cx))?;

            match self.inner.try_io(Interest::WRITABLE, || f(&sockref)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });
}

#[cfg_attr(docsrs, doc(auto_cfg = false))]
impl AsyncRead for TcpStream {
    #[inline]
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

#[cfg_attr(docsrs, doc(auto_cfg = false))]
impl AsyncWrite for TcpStream {
    #[inline]
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, buf)
    }

    #[inline]
    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[io::IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
    }

    #[inline]
    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }

    #[inline]
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    #[inline]
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

impl OwnedReadHalf {
    impl_tcp_stream_peek!(|self, buf| {
        // `tokio` requires the buffer to be initialized. Oh, performance be damned,
        // let's just zero it out.
        buf.fill(MaybeUninit::new(0));

        #[allow(unsafe_code, reason = "XXX")]
        // SAFETY: We just filled the buffer with zeroes, so it is now
        // initialized.
        let buf = unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len()) };

        self.inner.as_ref().peek(buf).await
    });

    impl_tcp_stream_poll_peek!(|self, cx, buf| {
        self.inner.as_ref().poll_peek(cx, &mut ReadBuf::uninit(buf))
    });

    impl_tcp_stream_read!(|self, buf| {
        loop {
            self.inner.readable().await?;

            match self.inner.try_read_buf(&mut ReadBuf::uninit(buf)) {
                read @ Ok(_) => break read,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                ret @ Err(_) => return ret,
            }
        }
    });

    impl_tcp_stream_read_fn!(|self, f| {
        let this = self.inner.as_ref();

        let sockref = SockRef::from(this);

        this.async_io(Interest::READABLE, move || f(&sockref))
    });

    impl_tcp_stream_poll_read_fn!(|self, cx, f| {
        let this = self.inner.as_ref();

        let sockref = SockRef::from(this);

        loop {
            ready!(this.poll_read_ready(cx))?;

            match this.try_io(Interest::READABLE, || f(&sockref)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });
}

#[cfg_attr(docsrs, doc(auto_cfg = false))]
impl AsyncRead for OwnedReadHalf {
    #[inline]
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl OwnedWriteHalf {
    impl_tcp_stream_write!(|self, buf| {
        loop {
            self.inner.writable().await?;

            match self.inner.try_write(buf) {
                write @ Ok(_) => break write,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                ret @ Err(_) => return ret,
            }
        }
    });

    impl_tcp_stream_write_fn!(|self, f| {
        let this = self.inner.as_ref();

        let sockref = SockRef::from(this);

        this.async_io(Interest::WRITABLE, move || f(&sockref))
    });

    impl_tcp_stream_poll_write_fn!(|self, cx, f| {
        let this = self.inner.as_ref();

        let sockref = SockRef::from(this);

        loop {
            ready!(this.poll_write_ready(cx))?;

            match this.try_io(Interest::WRITABLE, || f(&sockref)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });
}

#[cfg_attr(docsrs, doc(auto_cfg = false))]
impl AsyncWrite for OwnedWriteHalf {
    #[inline]
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, buf)
    }

    #[inline]
    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[io::IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
    }

    #[inline]
    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }

    #[inline]
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    #[inline]
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

impl TcpListener {
    impl_bound_socket_local_addr!(|self| {
        // (this comment is for formatting)
        self.inner.local_addr().map(From::from)
    });

    impl_tcp_listener_bind!(|addr| {
        let Some(addr) = addr.as_socket() else {
            return Err(io::Error::other("unsupported socket address family"));
        };

        tokio::net::TcpListener::bind(addr)
            .await
            .map(Self::from_inner)
    });

    impl_tcp_listener_accept!(|self| {
        loop {
            match self.inner.accept().await {
                Ok((stream, addr)) => {
                    break Ok((TcpStream::from_inner(stream), SockAddr::from(addr)));
                }
                Err(e)
                    if matches!(
                        e.kind(),
                        io::ErrorKind::ConnectionRefused
                            | io::ErrorKind::ConnectionAborted
                            | io::ErrorKind::ConnectionReset
                            | io::ErrorKind::BrokenPipe
                            | io::ErrorKind::Interrupted
                    ) =>
                {
                    // Connection errors are not deadly, just ignore and wait
                    // for the next connection.
                }
                Err(e) => break Err(e),
            }
        }
    });

    impl_tcp_listener_poll_accept!(|self, cx| {
        loop {
            match self.inner.poll_accept(cx) {
                Poll::Ready(Ok((stream, addr))) => {
                    break Poll::Ready(Ok((TcpStream::from_inner(stream), SockAddr::from(addr))));
                }
                Poll::Ready(Err(e))
                    if matches!(
                        e.kind(),
                        io::ErrorKind::ConnectionRefused
                            | io::ErrorKind::ConnectionAborted
                            | io::ErrorKind::ConnectionReset
                            | io::ErrorKind::BrokenPipe
                            | io::ErrorKind::Interrupted
                    ) =>
                {
                    // Connection errors are not deadly, just ignore and wait
                    // for the next connection.
                }
                Poll::Ready(Err(e)) => break Poll::Ready(Err(e)),
                Poll::Pending => break Poll::Pending,
            }
        }
    });
}

impl UdpSocket {
    impl_raw_socket_new!(|socket| {
        // (this comment is for formatting)
        tokio::net::UdpSocket::from_std(socket.into())
    });

    impl_raw_socket_bind!(|self, addr| {
        self.to_ref().bind(addr)?;

        Ok(self)
    });

    impl_bound_socket_local_addr!(|self| {
        // (this comment is for formatting)
        self.inner.local_addr().map(From::from)
    });

    impl_udp_socket_connect!(|self, addr| {
        // (this comment is for formatting)
        self.to_ref().connect(addr)
    });

    impl_connected_socket_peer_addr!(|self| {
        // (this comment is for formatting)
        self.inner.peer_addr().map(SockAddr::from)
    });

    impl_udp_socket_peek!(|self, buf| {
        // `tokio` requires the buffer to be initialized. Oh, performance be damned,
        // let's just zero it out.
        buf.fill(MaybeUninit::new(0));

        #[allow(unsafe_code, reason = "XXX")]
        // SAFETY: We just filled the buffer with zeroes, so it is now
        // initialized.
        let buf = unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len()) };

        self.inner.peek(buf).await
    });

    impl_udp_socket_poll_peek!(|self, cx, buf| {
        let mut buf = ReadBuf::uninit(buf);

        match self.inner.poll_peek(cx, &mut buf) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(buf.filled().len())),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    });

    impl_udp_socket_peek_sender!(|self| {
        // (this comment is for formatting)
        self.inner.peek_sender().await.map(SockAddr::from)
    });

    impl_udp_socket_poll_peek_sender!(|self, cx| {
        match self.inner.poll_peek_sender(cx) {
            Poll::Ready(Ok(addr)) => Poll::Ready(Ok(SockAddr::from(addr))),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    });

    impl_udp_socket_peek_from!(|self, buf| {
        // `tokio` requires the buffer to be initialized. Oh, performance be damned,
        // let's just zero it out.
        buf.fill(MaybeUninit::new(0));

        #[allow(unsafe_code, reason = "XXX")]
        // SAFETY: We just filled the buffer with zeroes, so it is now
        // initialized.
        let buf = unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len()) };

        self.inner
            .peek_from(buf)
            .await
            .map(|(peeked, addr)| (peeked, SockAddr::from(addr)))
    });

    impl_udp_socket_poll_peek_from!(|self, cx, buf| {
        let mut buf = ReadBuf::uninit(buf);

        match self.inner.poll_peek_from(cx, &mut buf) {
            Poll::Ready(Ok(addr)) => Poll::Ready(Ok((buf.filled().len(), SockAddr::from(addr)))),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    });

    impl_udp_socket_send!(|self, buf| {
        // (this comment is for formatting)
        self.inner.send(buf).await
    });

    impl_udp_socket_poll_send!(|self, cx, buf| {
        // (this comment is for formatting)
        self.inner.poll_send(cx, buf)
    });

    impl_udp_socket_send_to!(|self, buf, target| {
        let Some(addr) = target.as_socket() else {
            return Err(io::Error::other("unsupported socket address family"));
        };

        self.inner.send_to(buf, addr).await
    });

    impl_udp_socket_poll_send_to!(|self, cx, buf, target| {
        let Some(addr) = target.as_socket() else {
            return Poll::Ready(Err(io::Error::other("unsupported socket address family")));
        };

        self.inner.poll_send_to(cx, buf, addr)
    });

    impl_udp_socket_recv!(|self, buf| {
        // (this comment is for formatting)
        self.inner.recv_buf(&mut ReadBuf::uninit(buf)).await
    });

    impl_udp_socket_poll_recv!(|self, cx, buf| {
        let mut buf = ReadBuf::uninit(buf);

        match self.inner.poll_recv(cx, &mut buf) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(buf.filled().len())),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    });

    impl_udp_socket_recv_from!(|self, buf| {
        self.inner
            .recv_buf_from(&mut ReadBuf::uninit(buf))
            .await
            .map(|(read, addr)| (read, SockAddr::from(addr)))
    });

    impl_udp_socket_poll_recv_from!(|self, cx, buf| {
        let mut buf = ReadBuf::uninit(buf);

        match self.inner.poll_recv_from(cx, &mut buf) {
            Poll::Ready(Ok(addr)) => Poll::Ready(Ok((buf.filled().len(), SockAddr::from(addr)))),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    });

    impl_udp_socket_recv_fn!(|self, f| {
        let inner = SockRef::from(&self.inner);

        self.inner.async_io(Interest::READABLE, move || f(&inner))
    });

    impl_udp_socket_poll_recv_fn!(|self, cx, f| {
        let inner = SockRef::from(&self.inner);

        loop {
            ready!(self.inner.poll_recv_ready(cx))?;

            match self.inner.try_io(Interest::READABLE, || f(&inner)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });

    impl_udp_socket_send_fn!(|self, f| {
        let inner = SockRef::from(&self.inner);

        self.inner.async_io(Interest::WRITABLE, move || f(&inner))
    });

    impl_udp_socket_poll_send_fn!(|self, cx, f| {
        let inner = SockRef::from(&self.inner);

        loop {
            ready!(self.inner.poll_send_ready(cx))?;

            match self.inner.try_io(Interest::WRITABLE, || f(&inner)) {
                Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {}
                result => return Poll::Ready(result),
            }
        }
    });
}
