//! Utilities to cause interruptions in common traits, like Read/Write and Iterator.
use std::{
    io,
    sync::atomic::{AtomicBool, Ordering},
};

/// A wrapper for an inner iterator which will check for interruptions on each iteration, stopping the iteration when
/// that is requested.
pub struct Iter<'a, I> {
    /// The actual iterator to yield elements from.
    pub inner: I,
    should_interrupt: &'a AtomicBool,
}

impl<'a, I> Iter<'a, I>
where
    I: Iterator,
{
    /// Create a new iterator over `inner` which checks for interruptions on each iteration on `should_interrupt`.
    ///
    /// Note that this means the consumer of the iterator data should also be able to access `should_interrupt` and
    /// consider it when producing the final result to avoid claiming success even though the operation is only partially
    /// complete.
    pub fn new(inner: I, should_interrupt: &'a AtomicBool) -> Self {
        Iter {
            inner,
            should_interrupt,
        }
    }
}

impl<'a, I> Iterator for Iter<'a, I>
where
    I: Iterator,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.should_interrupt.load(Ordering::Relaxed) {
            return None;
        }
        self.inner.next()
    }
}

/// A wrapper for an inner iterator which will check for interruptions on each iteration.
pub struct IterWithErr<'a, I, EFN> {
    /// The actual iterator to yield elements from.
    pub inner: I,
    make_err: Option<EFN>,
    should_interrupt: &'a AtomicBool,
}

impl<'a, I, EFN, E> IterWithErr<'a, I, EFN>
where
    I: Iterator,
    EFN: FnOnce() -> E,
{
    /// Create a new iterator over `inner` which checks for interruptions on each iteration and calls `make_err()` to
    /// signal an interruption happened, causing no further items to be iterated from that point on.
    pub fn new(inner: I, make_err: EFN, should_interrupt: &'a AtomicBool) -> Self {
        IterWithErr {
            inner,
            make_err: Some(make_err),
            should_interrupt,
        }
    }
}

impl<'a, I, EFN, E> Iterator for IterWithErr<'a, I, EFN>
where
    I: Iterator,
    EFN: FnOnce() -> E,
{
    type Item = Result<I::Item, E>;

    fn next(&mut self) -> Option<Self::Item> {
        self.make_err.as_ref()?;
        if self.should_interrupt.load(Ordering::Relaxed) {
            return self.make_err.take().map(|f| Err(f()));
        }
        match self.inner.next() {
            Some(next) => Some(Ok(next)),
            None => {
                self.make_err = None;
                None
            }
        }
    }
}

/// A wrapper for implementers of [`std::io::Read`] or [`std::io::BufRead`] with interrupt support.
///
/// It fails a [read][std::io::Read::read] while an interrupt was requested.
pub struct Read<'a, R> {
    /// The actual implementor of [`std::io::Read`] to which interrupt support will be added.
    pub inner: R,
    /// The flag to trigger interruption
    pub should_interrupt: &'a AtomicBool,
}

impl<'a, R> io::Read for Read<'a, R>
where
    R: io::Read,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.should_interrupt.load(Ordering::Relaxed) {
            return Err(std::io::Error::new(std::io::ErrorKind::Other, "Interrupted"));
        }
        self.inner.read(buf)
    }
}

impl<'a, R> io::BufRead for Read<'a, R>
where
    R: io::BufRead,
{
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.inner.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.inner.consume(amt)
    }
}

/// A wrapper for implementers of [`std::io::Write`] with interrupt checks on each write call.
///
/// It fails a [write][std::io::Write::write] while an interrupt was requested.
pub struct Write<'a, W> {
    /// The actual implementor of [`std::io::Write`] to which interrupt support will be added.
    pub inner: W,
    /// The flag to trigger interruption
    pub should_interrupt: &'a AtomicBool,
}

impl<W> io::Write for Write<'_, W>
where
    W: std::io::Write,
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.should_interrupt.load(Ordering::Relaxed) {
            return Err(std::io::Error::new(std::io::ErrorKind::Other, "Interrupted"));
        }
        self.inner.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        // Don't interrupt here, allow flushes to happen to prefer disk consistency.
        self.inner.flush()
    }
}

impl<W> io::Seek for Write<'_, W>
where
    W: std::io::Seek,
{
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.inner.seek(pos)
    }
}