//! I/O extension: `splice(2)`.

cfg_select! {
    docsrs => {
        mod macros;
        mod pipe;
        pub mod traits;
    }
    _ => {
        mod macros;
        mod pipe;
        mod traits;
    }
}

use std::cmp::min;
use std::future::poll_fn;
use std::num::{NonZeroU64, NonZeroUsize};
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use std::{fmt, hint, io};

use rustix::pipe::{SpliceFlags, splice};

use self::macros::test_readiness;
use self::pipe::Pipe;
use self::traits::{AsyncFile, AsyncRead, AsyncStream, AsyncWrite};
use crate::util::unlikely;

#[must_use = "Futures do nothing unless you `.await` or poll them"]
/// Moves data from the source file descriptor to the destination file
/// descriptor without copying between kernel address space and user address
/// space, using [`splice(2)`].
///
/// [`splice(2)`]: https://man7.org/linux/man-pages/man2/splice.2.html
pub struct SpliceIO<L = (), R = ()> {
    ctx: Option<Ctx>,
    src: L,
    dst: R,
}

impl<L, R> fmt::Debug for SpliceIO<L, R> {
    #[cfg_attr(coverage_nightly, coverage(off))]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SpliceIO").finish_non_exhaustive()
    }
}

impl SpliceIO {
    /// Creates a new [`SpliceIO`] instance moving data from `src` to
    /// `dst`.
    pub fn new<L, R>(src: L, dst: R) -> io::Result<SpliceIO<L, R>>
    where
        L: AsyncRead,
        R: AsyncWrite,
    {
        Ok(SpliceIO {
            ctx: Some(Ctx {
                pipe: Pipe::new()?,
                off_in: Offset::NONE,
                off_out: Offset::NONE,
                target: u64::MAX,
                drained: 0,
                pumped: 0,
            }),
            src,
            dst,
        })
    }
}

impl<L, R> SpliceIO<L, R> {
    #[allow(clippy::missing_panics_doc, reason = "XXX")]
    /// Sets the total length of data to move.
    pub const fn with_len(mut self, target: u64) -> Self {
        self.ctx
            .as_mut()
            .expect("invalid state: ctx is None")
            .target = target;

        self
    }
}

impl<L, R> SpliceIO<L, R>
where
    L: AsyncFile,
{
    #[allow(clippy::missing_panics_doc, reason = "XXX")]
    /// Sets the read offset of the source file descriptor, which must be a
    /// regular file.
    pub const fn with_off_in(mut self, (off_in, target): (Offset, u64)) -> Self {
        let ctx = self.ctx.as_mut().expect("invalid state: ctx is None");
        ctx.off_in = off_in;
        ctx.target = target;

        self
    }
}

impl<L, R> SpliceIO<L, R>
where
    R: AsyncFile,
{
    #[allow(clippy::missing_panics_doc, reason = "XXX")]
    /// Sets the write offset of the destination file descriptor, which must be
    /// a regular file.
    pub const fn with_off_out(mut self, (off_out, target): (Offset, u64)) -> Self {
        let ctx = self.ctx.as_mut().expect("invalid state: ctx is None");
        ctx.off_out = off_out;
        ctx.target = target;

        self
    }
}

impl<L, R> SpliceIO<L, R>
where
    L: AsyncRead,
    R: AsyncWrite,
{
    /// Starts moving data from the source file descriptor to the destination
    /// file descriptor.
    pub async fn execute(mut self) -> SpliceIORet {
        poll_fn(|cx| self.poll_execute(cx)).await
    }

    #[allow(clippy::too_many_lines, reason = "XXX")]
    #[allow(clippy::missing_panics_doc, reason = "XXX")]
    /// The polling version of [`execute`].
    ///
    /// ## Notes
    ///
    /// - On multiple calls, only the `Waker` from the `Context` passed to the
    ///   most recent call is scheduled to receive a wakeup.
    ///
    /// [`execute`]: Self::execute
    pub fn poll_execute(&mut self, cx: &mut Context<'_>) -> Poll<SpliceIORet> {
        const MAX_SIZE_PER_SPLICE_SYSCALL: u64 = isize::MAX as u64;

        let Some(ctx) = self.ctx.as_mut() else {
            hint::cold_path();

            return Poll::Ready(
                self.shutdown(Some(io::Error::other("invalid state: poll after ready"))),
            );
        };

        let ret = loop {
            let Some(remaining) = ctx.target.checked_sub(ctx.pumped).and_then(NonZeroU64::new)
            else {
                break Ok(());
            };

            let Some(pumping) = ctx.drained.checked_sub(ctx.pumped) else {
                hint::cold_path();

                break Err(io::Error::other("unexpected state: drained < pumped"));
            };

            if pumping == 0 {
                #[allow(
                    clippy::cast_possible_truncation,
                    reason = "pumping <= MAX_SIZE_PER_SPLICE_SYSCALL (isize::MAX)"
                )]
                // Poll draining data from the source file descriptor to the
                // pipe.
                let draining = min(
                    ctx.pipe.capacity().get(),
                    min(remaining.get(), MAX_SIZE_PER_SPLICE_SYSCALL) as usize,
                );

                let fut = self.src.poll_read_fn(cx, |src| {
                    #[allow(unsafe_code, reason = "XXX")]
                    let off_in = unsafe { ctx.off_in.get() };

                    // ! It's guaranteed that the pipe is empty at this point and the target size is
                    // ! smaller than the pipe capacity. When EAGAIN is returned, we can safely
                    // ! assume that the source file descriptor is indeed not ready for reading.
                    splice(
                        src,
                        off_in,
                        ctx.pipe.wfd(),
                        None,
                        draining,
                        SpliceFlags::NONBLOCK,
                    )
                    .map(NonZeroUsize::new)
                    .map_err(|e| {
                        if e.raw_os_error() == libc::EAGAIN
                            && unlikely!(test_readiness!(&src, IN).unwrap_or_default())
                        {
                            io::Error::other(format!(
                                "unexpected: source fd is readable but `splice(2)` returns \
                                 EAGAIN, ctx: {ctx:?}",
                            ))
                        } else {
                            io::Error::from(e)
                        }
                    })
                });

                let drained = match ready!(fut) {
                    Ok(Some(drained)) => drained,
                    // * Read EOF
                    Ok(None) => break Ok(()),
                    Err(e) if matches!(e.kind(), io::ErrorKind::Interrupted) => continue,
                    Err(e) => break Err(e),
                };

                debug_assert!(drained.get() <= draining, "unexpected: drained > draining");

                ctx.drained += drained.get() as u64;
            } else {
                debug_assert!(
                    pumping <= MAX_SIZE_PER_SPLICE_SYSCALL,
                    "unexpected: pumping > MAX_SIZE_PER_SPLICE_SYSCALL"
                );

                #[allow(
                    clippy::cast_possible_truncation,
                    reason = "pumping <= MAX_SIZE_PER_SPLICE_SYSCALL (isize::MAX)"
                )]
                let pumping = pumping as usize;

                // Poll pumping data from the pipe to the destination file
                // descriptor.
                let fut = self.dst.poll_write_fn(cx, |dst| {
                    #[allow(unsafe_code, reason = "XXX")]
                    let off_out = unsafe { ctx.off_out.get() };

                    // ! It's guaranteed that the bytes we require are available in the pipe. When
                    // ! EAGAIN is returned, we can safely assume that the destination file
                    // ! descriptor is indeed not ready for writing.
                    splice(
                        ctx.pipe.rfd(),
                        None,
                        dst,
                        off_out,
                        pumping,
                        SpliceFlags::NONBLOCK,
                    )
                    .map(NonZeroUsize::new)
                    .map_err(|e| {
                        if e.raw_os_error() == libc::EAGAIN
                            && unlikely!(test_readiness!(&dst, OUT).unwrap_or_default())
                        {
                            io::Error::other(format!(
                                "unexpected: destination fd is writable but `splice(2)` returns \
                                 EAGAIN, ctx: {ctx:?}"
                            ))
                        } else {
                            io::Error::from(e)
                        }
                    })
                });

                let pumped = match ready!(fut) {
                    Ok(Some(pumped)) => pumped,
                    // * Write EOF
                    Ok(None) => break Err(io::ErrorKind::WriteZero.into()),
                    Err(e) if matches!(e.kind(), io::ErrorKind::Interrupted) => continue,
                    Err(e) => break Err(e),
                };

                debug_assert!(pumped.get() <= pumping, "unexpected: pumped > pumping");

                ctx.pumped += pumped.get() as u64;
            }
        };

        Poll::Ready(self.shutdown(ret.err()))
    }

    /// Cancels splicing from `L` to `R`.
    ///
    /// This is generally useful when one direction of a bidirectional splice
    /// operation has completed (e.g., due to EOF), and we don't want to wait
    /// for the other direction to complete, this is called a brutal shutdown.
    pub fn cancel(&mut self) -> SpliceIORet {
        self.shutdown(None)
    }

    fn shutdown(&mut self, err: Option<io::Error>) -> SpliceIORet {
        let Some(ctx) = self.ctx.take() else {
            return SpliceIORet::new(
                err.or_else(|| Some(io::Error::other("invalid state: cancel after ready"))),
                0,
                0,
            );
        };

        // TODO: recycle the pipe.
        // ...

        SpliceIORet::new(err, ctx.drained, ctx.pumped)
    }
}

impl<L, R> Future for SpliceIO<L, R>
where
    L: AsyncRead + Unpin,
    R: AsyncWrite + Unpin,
{
    type Output = SpliceIORet;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.get_mut().poll_execute(cx)
    }
}

#[non_exhaustive]
#[derive(Debug)]
/// The result of a completed [`SpliceIO`] execution.
pub struct SpliceIORet {
    /// The error if the operation fails, or `None` on success.
    pub error: Option<io::Error>,

    /// The total bytes drained from the source file descriptor when the splice
    /// operation finishes.
    pub drained: u64,

    /// The total bytes pumped to the destination file descriptor when the
    /// splice operation finishes.
    pub pumped: u64,
}

impl SpliceIORet {
    const fn new(error: Option<io::Error>, drained: u64, pumped: u64) -> Self {
        Self {
            error,
            drained,
            pumped,
        }
    }
}

#[must_use = "Futures do nothing unless you `.await` or poll them"]
/// The bidirectional version of [`SpliceIO`].
pub struct SpliceBidiIO<TxL = (), TxR = (), RxL = (), RxR = ()> {
    tx: SpliceIO<TxL, TxR>,
    rx: SpliceIO<RxL, RxR>,
}

#[cfg_attr(coverage_nightly, coverage(off))]
impl<TxL, TxR, RxL, RxR> fmt::Debug for SpliceBidiIO<TxL, TxR, RxL, RxR> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SpliceBidiIO").finish_non_exhaustive()
    }
}

impl SpliceBidiIO {
    #[allow(clippy::type_complexity, reason = "XXX")]
    /// Prepares to move data between `l` and `r` in both directions.
    pub fn new<L, R>(
        l: L,
        r: R,
    ) -> io::Result<SpliceBidiIO<L::ReadHalf, R::WriteHalf, R::ReadHalf, L::WriteHalf>>
    where
        L: AsyncStream,
        R: AsyncStream,
    {
        let (l_r, l_w) = l.split();
        let (r_r, r_w) = r.split();

        Ok(SpliceBidiIO {
            tx: SpliceIO::new(l_r, r_w)?,
            rx: SpliceIO::new(r_r, l_w)?,
        })
    }
}

impl<TxL, TxR, RxL, RxR> SpliceBidiIO<TxL, TxR, RxL, RxR> {
    #[allow(clippy::missing_panics_doc, reason = "XXX")]
    /// Sets the total length of data to move in each direction.
    pub const fn with_len(mut self, target: u64) -> Self {
        self.tx
            .ctx
            .as_mut()
            .expect("invalid state: ctx is None")
            .target = target;
        self.rx
            .ctx
            .as_mut()
            .expect("invalid state: ctx is None")
            .target = target;
        self
    }
}

impl<TxL, TxR, RxL, RxR> SpliceBidiIO<TxL, TxR, RxL, RxR>
where
    TxL: AsyncRead,
    TxR: AsyncWrite,
    RxL: AsyncRead,
    RxR: AsyncWrite,
{
    /// Starts moving data between the two file descriptors in both directions.
    pub async fn execute(mut self) -> SpliceBidiIORet {
        poll_fn(|cx| self.poll_execute(cx)).await
    }

    /// The polling version of [`execute`].
    ///
    /// ## Notes
    ///
    /// - On multiple calls, only the `Waker` from the `Context` passed to the
    ///   most recent call is scheduled to receive a wakeup.
    ///
    /// [`execute`]: Self::execute
    pub fn poll_execute(&mut self, cx: &mut Context<'_>) -> Poll<SpliceBidiIORet> {
        let tx = self.tx.poll_execute(cx);
        let rx = self.rx.poll_execute(cx);

        match (tx, rx) {
            (Poll::Pending, Poll::Pending) => Poll::Pending,
            (Poll::Ready(tx), Poll::Pending) => {
                // One direction has completed, we cancel the other direction and return the
                // result immediately, i.e., brutal shutdown.
                Poll::Ready(SpliceBidiIORet {
                    tx,
                    rx: self.rx.cancel(),
                })
            }
            (Poll::Pending, Poll::Ready(rx)) => {
                // One direction has completed, we cancel the other direction and return the
                // result immediately, i.e., brutal shutdown.
                Poll::Ready(SpliceBidiIORet {
                    tx: self.tx.cancel(),
                    rx,
                })
            }
            (Poll::Ready(tx), Poll::Ready(rx)) => {
                // Unlikely, though.
                Poll::Ready(SpliceBidiIORet { tx, rx })
            }
        }
    }
}

impl<TxL, TxR, RxL, RxR> Future for SpliceBidiIO<TxL, TxR, RxL, RxR>
where
    TxL: AsyncRead + Unpin,
    TxR: AsyncWrite + Unpin,
    RxL: AsyncRead + Unpin,
    RxR: AsyncWrite + Unpin,
{
    type Output = SpliceBidiIORet;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.get_mut().poll_execute(cx)
    }
}

#[non_exhaustive]
#[derive(Debug)]
/// The result of a completed [`SpliceBidiIO`] execution.
pub struct SpliceBidiIORet {
    /// TX (from L to R).
    pub tx: SpliceIORet,
    /// RX (from R to L).
    pub rx: SpliceIORet,
}

#[derive(Debug)]
/// Runtime state for a [`SpliceIO`] operation.
struct Ctx {
    /// The pipe used when moving data between two file descriptors.
    ///
    /// Although `splice(2)` is nominally "zero-copy", the Linux kernel requires
    /// a pipe to serve as an "intermediary". There is overhead in creating
    /// and destroying pipes, and whether there is a reasonable way
    /// to implement a pipe pool remains to be explored.
    pipe: Pipe,

    /// The target bytes of data to move.
    target: u64,

    /// The (current) read offset of the source file descriptor.
    off_in: Offset,

    /// The (current) write offset of the destination file descriptor.
    off_out: Offset,

    /// The total bytes drained from the source file descriptor to the pipe.
    drained: u64,

    /// The total bytes pumped from the pipe to the destination file descriptor.
    pumped: u64,
}

wrapper_lite::wrapper!(
    #[repr(transparent)]
    #[derive(Debug)]
    /// A read or write offset for the source or destination file descriptor.
    pub struct Offset(Option<NonZeroU64>);
);

impl Offset {
    const NONE: Self = Self::from_inner(None);

    /// Creates a new [`Offset`] for reading from the source file descriptor or
    /// writing to the destination file descriptor.
    ///
    /// The file descriptor must be a regular file.
    ///
    /// ## Parameters
    ///
    /// - `len`
    ///
    ///   The total length of the file.
    ///
    /// - `from`
    ///
    ///   The start offset of the file, or `None` to start from the
    ///   beginning.
    ///
    /// - `to`
    ///
    ///   The end offset of the file, or `None` to end at the length of
    ///   the file.
    ///
    /// ## Returns
    ///
    /// The offset, and the total length of the data to move.
    pub fn new(len: u64, from: Option<u64>, to: Option<u64>) -> io::Result<(Self, u64)> {
        let offset = from.and_then(NonZeroU64::new);
        let len = match (offset, to.and_then(NonZeroU64::new)) {
            (Some(offset), Some(end)) => {
                if offset > end || end.get() > len {
                    return Err(io::Error::other("invalid range"));
                }

                end.get() - offset.get()
            }
            (Some(offset), None) => {
                if offset.get() > len {
                    return Err(io::Error::other("invalid `from`"));
                }

                len - offset.get()
            }
            (None, Some(end)) => {
                if end.get() > len {
                    return Err(io::Error::other("invalid `to`"));
                }
                end.get()
            }
            (None, None) => len,
        };

        Ok((Self::from_inner(offset), len))
    }

    #[inline]
    #[allow(unsafe_code, reason = "XXX")]
    /// ## Safety
    ///
    /// The returned `&mut u64` is actually a `&mut NonZeroU64`, and must
    /// never be modified to zero.
    const unsafe fn get(&mut self) -> Option<&mut u64> {
        #[allow(unsafe_code, reason = "XXX")]
        // SAFETY: the caller must ensure that the returned `&mut u64` is never modified to zero.
        match &mut self.inner {
            Some(off) => unsafe { Some(&mut *(&raw mut *off).cast()) },
            _ => None,
        }
    }
}
