/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2023 metamuffin <metamuffin.org>
*/
use crate::{matroska::MatroskaTag, Master};
use anyhow::Result;

pub trait IterWithPos {
    type Item;
    fn next(&mut self) -> Option<Self::Item>;
    fn position(&self) -> usize;
}

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

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

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

    pub fn position(&self) -> usize {
        self.inner.position()
    }

    pub fn n(&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 self.n().is_some() {}
    }
}