use futures::{Stream, StreamExt};

use crate::{
    datasets::DatasetView, graphs::default::DefaultFilter, impl_default_dataset_view, prelude::DatasetError, quads::{filter::QuadFilter, Quad}
};

pub struct FilterStream<S, F> {
    stream: S,
    filter: F,
}

#[macro_export]
macro_rules! impl_default_dataset_filter {
    () => {
        type Filter<FilterType>
            = $crate::graphs::default::DefaultFilter<Self, FilterType>
        where
            FilterType: $crate::quads::filter::QuadFilter + ::core::marker::Unpin;

        async fn filter<FilterType>(
            self,
            filter: FilterType,
        ) -> Result<Self::Filter<FilterType>, $crate::datasets::DatasetError>
        where
            FilterType: $crate::quads::filter::QuadFilter + ::core::marker::Unpin,
        {
            Ok($crate::graphs::default::DefaultFilter::new(self, filter))
        }
    };
}

impl<D, F> DatasetView for DefaultFilter<D, F>
where
    D: DatasetView,
    F: QuadFilter + Unpin,
{
    type Stream = FilterStream<D::Stream, F>;

    async fn stream(self) -> Result<Self::Stream, crate::prelude::DatasetError> {
        Ok(Self::Stream {
            stream: self.view.stream().await?,
            filter: self.filter,
        })
    }

    impl_default_dataset_view!();
}

impl<S, F> Stream for FilterStream<S, F>
where
    S: Stream<Item = Result<Quad, DatasetError>> + Unpin,
    F: QuadFilter + Unpin,
{
    type Item = Result<Quad, DatasetError>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        loop {
            match self.stream.poll_next_unpin(cx) {
                core::task::Poll::Ready(Some(Ok(q))) => {
                    if self.filter.filter(q.as_quad_ref()) {
                        break core::task::Poll::Ready(Some(Ok(q)));
                    }
                }
                err @ core::task::Poll::Ready(Some(Err(_))) => break err,
                end @ core::task::Poll::Ready(None) => break end,
                core::task::Poll::Pending => break core::task::Poll::Pending,
            }
        }
    }
}
