use super::Document;
use super::Element;
use super::Heading;
use super::Object;
use super::Section;
use crate::parser::DocumentElement;

pub enum Token<'r, 's> {
    Document(&'r Document<'s>),
    Heading(&'r Heading<'s>),
    Section(&'r Section<'s>),
    Object(&'r Object<'s>),
    Element(&'r Element<'s>),
}

impl<'r, 's> Token<'r, 's> {
    pub fn iter_tokens(&self) -> Box<dyn Iterator<Item = Token<'r, 's>> + '_> {
        match self {
            Token::Document(document) => Box::new(
                document
                    .zeroth_section
                    .iter()
                    .map(Token::Section)
                    .chain(document.children.iter().map(Token::Heading)),
            ),
            Token::Heading(heading) => Box::new(heading.title.iter().map(Token::Object).chain(
                heading.children.iter().map(|de| match de {
                    DocumentElement::Heading(ref obj) => Token::Heading(obj),
                    DocumentElement::Section(ref obj) => Token::Section(obj),
                }),
            )),
            Token::Section(section) => Box::new(section.children.iter().map(Token::Element)),
            Token::Object(_) => panic!(),
            Token::Element(_) => panic!(),
        }
    }
}