use async_trait::async_trait;
use tokio::io::{AsyncBufRead, AsyncBufReadExt};

#[async_trait]
pub trait AsyncPeek {
    async fn peek(&mut self, buf: &mut [u8]) -> tokio::io::Result<usize>;
}

#[async_trait]
impl<T> AsyncPeek for T
where
    T: AsyncBufRead + Unpin + Send,
{
    async fn peek(&mut self, buf: &mut [u8]) -> tokio::io::Result<usize> {
        let filled = self.fill_buf().await?;
        if filled.len() >= buf.len() {
            buf.copy_from_slice(&filled[..buf.len()]);
            Ok(buf.len())
        } else {
            buf[..filled.len()].copy_from_slice(filled);
            Ok(filled.len())
        }
    }
}