use crate::{matroska::MatroskaTag, Master};
use anyhow::Result;

pub struct Unflat<'a> {
    pub item: MatroskaTag,
    pub children: Option<Unflatten<'a>>,
}

pub struct Unflatten<'a> {
    inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>,
    stop: bool,
    end: Option<MatroskaTag>,
}

impl<'a> Unflatten<'a> {
    pub fn new(inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>) -> Self {
        Self {
            inner,
            stop: false,
            end: None,
        }
    }
    pub fn new_with_end(
        inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>,
        start: MatroskaTag,
    ) -> Self {
        Self {
            inner,
            stop: false,
            end: Some(MatroskaTag::construct_master(start.id(), Master::End).unwrap()),
        }
    }

    pub fn next(&mut self) -> Option<Result<Unflat>> {
        if self.stop {
            return None;
        }
        match self.inner.next() {
            None => None,
            Some(Err(e)) => Some(Err(e)),
            Some(Ok(item)) => {
                let master = MatroskaTag::is_master(item.id()).unwrap();
                if Some(&item) == self.end.as_ref() {
                    self.stop = true;
                    None
                } else {
                    Some(Ok(Unflat {
                        children: if master {
                            let end =
                                MatroskaTag::construct_master(item.id(), Master::End).unwrap();
                            if end == item {
                                return None;
                            }
                            Some(Unflatten {
                                inner: self.inner,
                                stop: false,
                                end: Some(end),
                            })
                        } else {
                            None
                        },
                        item,
                    }))
                }
            }
        }
    }
}

impl Drop for Unflatten<'_> {
    fn drop(&mut self) {
        while let Some(_) = self.next() {}
    }
}