use crate::MarkdownExt;
use equt_md_frontmatter::FrontMatter;
use std::cell::{Ref, RefCell};
use std::rc::Rc;

enum HeadState {
    Head,
    Origin,
}

pub struct Head<I, H, F> {
    frontmatter: Rc<RefCell<Option<FrontMatter>>>,
    head: Option<H>,
    origin: I,
    f: F,
    state: HeadState,
}

impl<I, G, H, F, T> Head<I, H, F>
where
    G: IntoIterator<Item = T, IntoIter = H>,
    H: Iterator<Item = T>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
{
    pub(crate) fn new(frontmatter: Rc<RefCell<Option<FrontMatter>>>, origin: I, f: F) -> Self {
        Self {
            frontmatter,
            head: None,
            origin,
            f,
            state: HeadState::Head,
        }
    }
}

impl<I, H, F, G, T> Iterator for Head<I, H, F>
where
    I: Iterator<Item = T>,
    H: Iterator<Item = T>,
    G: IntoIterator<Item = T, IntoIter = H>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.head.is_none() {
            self.head = Some((self.f)(self.frontmatter.borrow()).into_iter());
        }
        match self.state {
            HeadState::Head => match self.head.as_mut().unwrap().next() {
                node @ Some(_) => node,
                None => {
                    self.state = HeadState::Origin;
                    self.next()
                }
            },
            HeadState::Origin => self.origin.next(),
        }
    }
}

impl<I, H, F, G, T> MarkdownExt<T> for Head<I, H, F>
where
    I: Iterator<Item = T>,
    H: Iterator<Item = T>,
    G: IntoIterator<Item = T, IntoIter = H>,
    F: Fn(Ref<Option<FrontMatter>>) -> G,
{
    fn frontmatter(&self) -> Rc<RefCell<Option<FrontMatter>>> {
        self.frontmatter.clone()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn normal() {
        assert_eq!(
            Head::new(Rc::new(RefCell::new(None)), 3..6, |_| 1..3).collect::<Vec<_>>(),
            vec![1, 2, 3, 4, 5]
        );
    }
}