use std::{io, fmt, cmp};
use std::io::BufRead;

use BufReadGrow;

/// An adapter that retries reading/writing operations of the underlying reader or writer.
///
/// This struct is generally created by calling `retry()` on a reader or writer.
/// Please see the documentation of [`Read::retry`] and [`Write::retry`] for more details.
///
/// [`Read::retry`]: ./trait.Read.html#method.retry
/// [`Write::retry`]: ./trait.Write.html#method.retry
pub struct Retry<I> {
    inner: I,
}

impl<I> Retry<I> {
    #[inline]
    pub fn new(inner: I) -> Retry<I> {
        Retry { inner: inner }
    }
}

impl<I: io::Read> io::Read for Retry<I> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        loop {
            match self.inner.read(buf) {
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                other => return other,
            }
        }
    }

    #[inline]
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        self.inner.read_to_end(buf)
    }

    #[inline]
    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        self.inner.read_to_string(buf)
    }

    #[inline]
    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.inner.read_exact(buf)
    }
}

impl<I: io::Write> io::Write for Retry<I> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        loop {
            match self.inner.write(buf) {
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                other => return other,
            }
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        loop {
            match self.inner.flush() {
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                other => return other,
            }
        }
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.inner.write_all(buf)
    }

    #[inline]
    fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()> {
        self.inner.write_fmt(fmt)
    }
}

impl<R: io::BufRead> io::BufRead for Retry<R> {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        while let Err(e) = self.inner.fill_buf() {
            if e.kind() != io::ErrorKind::Interrupted {
                return Err(e);
            }
        }
        self.inner.fill_buf()
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        self.inner.consume(amt)
    }

    #[inline]
    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
        self.inner.read_until(byte, buf)
    }

    #[inline]
    fn read_line(&mut self, buf: &mut String) -> io::Result<usize> {
        self.inner.read_line(buf)
    }
}

impl<R: BufReadGrow> BufReadGrow for Retry<R> {
    #[inline]
    fn grow_buf(&mut self) -> io::Result<&[u8]> {
        while let Err(e) = self.inner.grow_buf() {
            if e.kind() != io::ErrorKind::Interrupted {
                return Err(e);
            }
        }
        self.inner.fill_buf()
    }
}

/// An adapter that restarts from the beginning after EOF is reached.
///
/// This struct is generally created by calling `repeat()` on a reader.
/// Please see the documentation of [`Read::repeat`] for more details.
///
/// [`Read::repeat`]: ./trait.Read.html#method.repeat
pub struct Repeat<R> {
    inner: R,
}

impl<I> Repeat<I> {
    #[inline]
    pub fn new(inner: I) -> Repeat<I> {
        Repeat { inner: inner }
    }
}

impl<I: io::Read + io::Seek> io::Read for Repeat<I> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner.read(buf) {
            Ok(0) => {
                try!(self.inner.seek(io::SeekFrom::Start(0)));
                self.inner.read(buf)
            }
            Ok(n) => Ok(n),
            Err(e) => Err(e),
        }
    }
}

impl<I: io::BufRead + io::Seek> io::BufRead for Repeat<I> {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if try!(self.inner.fill_buf()).is_empty() {
            try!(self.inner.seek(io::SeekFrom::Start(0)));
        }
        self.inner.fill_buf()
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        self.inner.consume(amt)
    }
}

/// Adapter which limits the bytes read from / written to an underlying reader / writer.
///
/// This struct is generally created by calling `take()` on a reader/writer.
/// Please see the documentation of [`Read::take`] and [`Write::take`] for more details.
///
/// [`Read::take`]: ./trait.Read.html#method.take
/// [`Write::take`]: ./trait.Write.html#method.take
pub struct Take<T> {
    inner: T,
    limit: u64,
}

impl<T> Take<T> {
    #[inline]
    pub fn new(inner: T, limit: u64) -> Take<T> {
        Take {
            inner: inner,
            limit: limit,
        }
    }

    /// Returns the number of bytes that can be read before this instance will return EOF.
    ///
    /// Note
    /// ====
    /// This instance may reach EOF after reading fewer bytes than indicated by
    /// this method if the underlying `Read` instance reaches EOF.
    #[inline]
    pub fn limit(&self) -> u64 {
        self.limit
    }
}

impl<T: io::Read> io::Read for Take<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        // Don't call into inner reader at all at EOF because it may still block
        if self.limit == 0 {
            return Ok(0);
        }

        let max = cmp::min(buf.len() as u64, self.limit) as usize;
        let n = try!(self.inner.read(&mut buf[..max]));
        self.limit -= n as u64;
        Ok(n)
    }
}

impl<T: io::BufRead> io::BufRead for Take<T> {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        // Don't call into inner reader at all at EOF because it may still block
        if self.limit == 0 {
            return Ok(&[]);
        }

        let buf = try!(self.inner.fill_buf());
        let cap = cmp::min(buf.len() as u64, self.limit) as usize;
        Ok(&buf[..cap])
    }

    #[inline]
    fn consume(&mut self, amt: usize) {
        // Don't let callers reset the limit by passing an overlarge value
        let amt = cmp::min(amt as u64, self.limit) as usize;
        self.limit -= amt as u64;
        self.inner.consume(amt);
    }
}

impl<T: BufReadGrow> BufReadGrow for Take<T> {
    fn grow_buf(&mut self) -> io::Result<&[u8]> {
        // Don't call into inner reader at all at EOF because it may still block
        if self.limit == 0 || self.limit == try!(self.fill_buf()).len() as u64 {
            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Stream is already at EOF"));
        }

        let buf = try!(self.inner.grow_buf());
        let cap = cmp::min(buf.len() as u64, self.limit) as usize;
        Ok(&buf[..cap])
    }
}

impl<T: io::Write> io::Write for Take<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.limit == 0 {
            return Ok(0);
        }

        let amt = cmp::min(self.limit, buf.len() as u64) as usize;
        let amt = try!(self.inner.write(&buf[..amt]));
        self.limit -= amt as u64;
        Ok(amt)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use Stream;

    #[test]
    fn take_write() {
        let mut c = io::Cursor::new(&b"0123456789"[..]);
        let mut v = Vec::new();
        assert_eq!(::copy(&mut c, &mut v.by_ref().take(0)).unwrap_err().kind(),
                   io::ErrorKind::WriteZero);
        assert_eq!(v, b"");

        let mut c = io::Cursor::new(&b"0123456789"[..]);
        let mut v = Vec::new();
        assert_eq!(::copy(&mut c, &mut v.by_ref().take(9)).unwrap_err().kind(),
                   io::ErrorKind::WriteZero);
        assert_eq!(v, b"012345678");

        let mut c = io::Cursor::new(&b"0123456789"[..]);
        let mut v = Vec::new();
        ::copy(&mut c, &mut v.by_ref().take(10)).unwrap();
        assert_eq!(v, b"0123456789");

        let mut c = io::Cursor::new(&b"0123456789"[..]);
        let mut v = Vec::new();
        ::copy(&mut c, &mut v.by_ref().take(11)).unwrap();
        assert_eq!(v, b"0123456789");
    }
}