use futures_core::{ready, Stream};
use pin_project::{pin_project, project};
use std::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

/// A `ForEach`-like future.
#[must_use = "futures do nothing unless polled"]
#[pin_project]
pub struct StreamProcessor<S, P, F>
where
    F: Future,
{
    #[pin]
    receiver: S,
    processor: P,
    #[pin]
    state: State<F>,
}

impl<S, P, F> StreamProcessor<S, P, F>
where
    F: Future,
{
    pub fn new(receiver: S, processor: P) -> Self {
        StreamProcessor {
            receiver,
            processor,
            state: State::WaitingForItem,
        }
    }
}

#[pin_project]
enum State<F> {
    WaitingForItem,
    RunningProcessor(#[pin] F),
}

impl<S, P, F> Future for StreamProcessor<S, P, F>
where
    S: Stream,
    P: FnMut(S::Item) -> F,
    F: Future,
{
    type Output = ();

    #[project]
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            let this = self.as_mut().project();

            #[project]
            let next = match this.state.project() {
                State::WaitingForItem => {
                    if let Some(item) = ready!(this.receiver.poll_next(ctx)) {
                        let fut = (this.processor)(item);
                        State::RunningProcessor(fut)
                    } else {
                        return Poll::Ready(());
                    }
                }
                State::RunningProcessor(fut) => {
                    ready!(fut.poll(ctx));
                    State::WaitingForItem
                }
            };

            self.as_mut().project().state.set(next);
        }
    }
}