use std::fmt;
use std::pin::Pin;
use std::task::{self, Poll};

use futures::prelude::*;
use futures::stream::{FuturesOrdered, Stream, StreamFuture};

/// [`SelectAll`](futures::stream::SelectAll) but ordered
#[pin_project::pin_project]
pub struct ConcatAll<St: Stream + Unpin> {
    #[pin]
    inner: FuturesOrdered<StreamFuture<St>>,
}

impl<St: Stream + Unpin> ConcatAll<St> {
    pub fn new() -> Self {
        Self {
            inner: FuturesOrdered::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

    pub fn push(&mut self, stream: St) {
        use futures::StreamExt;

        self.inner.push(stream.into_future());
    }
}

impl<St: Stream + Unpin> fmt::Debug for ConcatAll<St> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ConcatAll").finish()
    }
}

impl<St: Stream + Unpin> Stream for ConcatAll<St> {
    type Item = St::Item;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            match futures::ready!(self.inner.poll_next_unpin(cx)) {
                Some((Some(item), remaining)) => {
                    self.push(remaining);
                    return Poll::Ready(Some(item));
                }
                Some((None, _)) => {}
                _ => return Poll::Ready(None),
            }
        }
    }
}