use super::{MarkdownExt, Share, FrontMatter};
use std::cell::{Ref, RefCell};

enum AfterState {
    Origin,
    Pseudo,
}

/// An iterator places events after the first occurrence of certain event.
///
/// This `struct` is created by [`after`] method on [`MarkdownExt`].
///
/// [`after`]: trait.MarkdownExt.html#method.after
/// [`MarkdownExt`]: trait.MarkdownExt.html
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[derive(new)]
pub struct After<I, S, G, F, P, T>
where
    I: Iterator<Item = T>,
    S: Iterator<Item = T>,
    G: IntoIterator<Item = T, IntoIter = S>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
    P: Fn(&T) -> bool,
{
    frontmatter: Share<RefCell<Option<FrontMatter>>>,
    iter: I,
    after: P,

    // Create pseudo iter
    f: F,

    // Delay its construction until needed
    #[new(default)]
    pseudo: Option<S>,
    #[new(value = "AfterState::Origin")]
    state: AfterState,
}

impl<I, S, G, F, P, T> MarkdownExt<T> for After<I, S, G, F, P, T>
where
    I: Iterator<Item = T>,
    S: Iterator<Item = T>,
    G: IntoIterator<Item = T, IntoIter = S>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
    P: Fn(&T) -> bool,
{
    fn frontmatter(&mut self) -> &mut Share<RefCell<Option<FrontMatter>>> {
        &mut self.frontmatter
    }
}

impl<I, S, G, F, P, T> Iterator for After<I, S, G, F, P, T>
where
    I: Iterator<Item = T>,
    S: Iterator<Item = T>,
    G: IntoIterator<Item = T, IntoIter = S>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
    P: Fn(&T) -> bool,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        match self.state {
            AfterState::Origin => match self.iter.next() {
                // If current item match the predication, return it and switch to
                // pseudo iterator.
                Some(t) if (self.after)(&t) => {
                    self.state = AfterState::Pseudo;
                    Some(t)
                }
                otherwise => otherwise,
            },
            // We need to check if the pseudo iterator has been created or not.
            AfterState::Pseudo => match self.pseudo.as_mut() {
                Some(iter) => match iter.next() {
                    node @ Some(_) => node,
                    // The whole pseudo iterator has been consumed, hand back to
                    // the original iterator.
                    None => {
                        self.state = AfterState::Origin;
                        self.iter.next()
                    }
                },
                None => {
                    self.pseudo = Some(
                        (self.f)(self.frontmatter.upgrade().unwrap().borrow()).into_iter(),
                    );

                    // Since we've written something into self.pseudo, it would definitely
                    // go to the Some arm above
                    self.next()
                }
            },
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::iter::{empty, once};
    use std::rc::Rc;

    #[test]
    fn only_once() {
        assert_eq!(
            After::new(
                Rc::new(RefCell::new(None)).into(),
                1..5,
                |&n| n % 2 == 0,
                |_| once(0)
            )
            .collect::<Vec<_>>(),
            vec![1, 2, 0, 3, 4]
        );
    }

    #[test]
    fn empty_pseudo() {
        assert_eq!(
            After::new(
                Rc::new(RefCell::new(None)).into(),
                1..5,
                |&n| n % 2 == 0,
                |_| empty()
            )
            .collect::<Vec<_>>(),
            vec![1, 2, 3, 4]
        );
    }

    #[test]
    fn never_match() {
        assert_eq!(
            After::new(
                Rc::new(RefCell::new(None)).into(),
                1..5,
                |&n| n % 2 == 5,
                |_| once(0)
            )
            .collect::<Vec<_>>(),
            vec![1, 2, 3, 4]
        );
    }

    #[test]
    fn lazy() {
        assert!(After::new(
                Rc::new(RefCell::new(None)).into(),
            1..5,
            |&n| n % 10 == 0,
            |_| once(0)
        )
        .pseudo
        .is_none());
    }
}
