mod audio;
mod code;
mod figure;
mod image;
mod label;
mod list;
mod margins;
mod math;
mod table;
mod video;

pub use audio::Audio;
pub use code::Code;
pub use figure::Figure;
pub use image::Image;
pub use label::{InlineWidget, InlineWidgetType, Label, Paragraph, TagType, TextTag};
pub use list::{ListBullet, ListItem, ListType};
pub use margins::Margins;
pub use math::{Math, MathLang};
use serde::{Deserialize, Serialize};
pub use table::{Table, TableCell, TableRow, TableRowType};
pub use video::Video;

#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Serialize)]
pub enum Element {
    Paragraph(Paragraph),
    Label(Label),
    Image(Image),
    Gallery(Vec<Image>),
    Audio(Audio),
    Video(Video),
    List(Vec<ListItem>),
    Table(Table),
    Container(Vec<Element>),
    Quote(Box<Element>),
    Figure(Figure),
    Math(Math),
    Code(Code),
}

impl Default for Element {
    fn default() -> Self {
        Self::Label(Label::default())
    }
}

impl Element {
    pub fn find_inline_math(&mut self, inline_math_delimiter: &str) {
        match self {
            Element::Paragraph(paragraph) => {
                paragraph.convert_inline_math(inline_math_delimiter);
            }
            Element::List(list_items) => {
                for list_item in list_items {
                    list_item.element.find_inline_math(inline_math_delimiter);
                }
            }
            Element::Table(table) => {
                for row in &mut table.rows {
                    for cell in &mut row.cells {
                        cell.element.find_inline_math(inline_math_delimiter);
                    }
                }
            }
            Element::Container(elements) => {
                for element in elements {
                    element.find_inline_math(inline_math_delimiter);
                }
            }
            Element::Quote(element) => element.find_inline_math(inline_math_delimiter),
            Element::Figure(figure) => figure.element.find_inline_math(inline_math_delimiter),
            _ => {}
        }
    }

    pub fn flatten(self) -> Option<Element> {
        if let Element::Container(boxed_elements) = self {
            let mut elements = Vec::new();

            for element in boxed_elements {
                if let Some(flattened) = element.flatten() {
                    match flattened {
                        Element::Container(mut nested_boxed_elements) => {
                            elements.append(&mut nested_boxed_elements)
                        }
                        _ => elements.push(flattened),
                    }
                }
            }

            Self::first_or_box(elements)
                .map(Element::paragraph_to_label)
                .map(Element::group_images)
        } else if let Element::Quote(quoted_element) = self {
            quoted_element
                .flatten()
                .map(Element::paragraph_to_label)
                .map(Element::group_images)
                .map(|e| Element::Quote(Box::new(e)))
        } else if let Element::Figure(Figure { element, caption }) = self {
            element
                .flatten()
                .map(Element::paragraph_to_label)
                .map(Element::group_images)
                .map(move |e| {
                    Element::Figure(Figure {
                        element: Box::new(e),
                        caption,
                    })
                })
        } else if let Element::List(list) = self {
            let list = list
                .into_iter()
                .map(|list_item| ListItem {
                    bullet: list_item.bullet,
                    element: list_item.element.paragraph_to_label().group_images(),
                })
                .collect();
            Some(Element::List(list))
        } else {
            Some(self)
        }
    }

    fn paragraph_to_label(self) -> Element {
        if let Element::Container(boxed_elements) = self {
            Element::Container(
                boxed_elements
                    .into_iter()
                    .map(Self::paragraph_to_label)
                    .collect(),
            )
        } else if let Element::Quote(quoted_element) = self {
            Element::Quote(Box::new(quoted_element.paragraph_to_label()))
        } else if let Element::Figure(Figure { element, caption }) = self {
            Element::Figure(Figure {
                element: Box::new(element.paragraph_to_label()),
                caption: caption.map(|c| Box::new(c.paragraph_to_label())),
            })
        } else if let Element::Table(table) = self {
            Element::Table(Table {
                caption: table.caption,
                width: table.width,
                rows: table
                    .rows
                    .into_iter()
                    .map(|row| TableRow {
                        row_type: row.row_type,
                        cells: row
                            .cells
                            .into_iter()
                            .map(|cell| TableCell {
                                width: cell.width,
                                height: cell.height,
                                element: cell.element.paragraph_to_label(),
                                is_header: cell.is_header,
                            })
                            .collect(),
                    })
                    .collect(),
            })
        } else if let Element::List(list) = self {
            let list = list
                .into_iter()
                .map(|list_item| ListItem {
                    bullet: list_item.bullet,
                    element: list_item.element.paragraph_to_label(),
                })
                .collect();
            Element::List(list)
        } else if let Element::Paragraph(paragraph) = self {
            Element::Label(Label::from(paragraph))
        } else {
            self
        }
    }

    pub fn group_paragraphs(elements: Vec<Element>) -> Vec<Element> {
        let mut grouped = Vec::new();

        for mut element in elements.into_iter() {
            if let Element::Paragraph(p) = &mut element
                && let Some(Element::Paragraph(last_p)) = grouped.last_mut()
            {
                let mut offset = last_p.text.chars().count() as i32;
                let ends_with_whitespace = last_p
                    .text
                    .chars()
                    .last()
                    .map(char::is_whitespace)
                    .unwrap_or(false);
                let starts_with_whitespace = p
                    .text
                    .chars()
                    .next()
                    .map(char::is_whitespace)
                    .unwrap_or(false);
                if !ends_with_whitespace && !starts_with_whitespace {
                    last_p.text.push(' ');
                    offset += 1;
                }
                let mut offset_tags = p
                    .tags
                    .drain(..)
                    .map(|tag| TextTag {
                        start: tag.start + offset,
                        stop: tag.stop + offset,
                        tag_type: tag.tag_type,
                    })
                    .collect();
                let mut offset_widgets = p
                    .widgets
                    .drain(..)
                    .map(|widget| InlineWidget {
                        offset: widget.offset + offset,
                        widget_type: widget.widget_type,
                    })
                    .collect();

                last_p.ids.extend(p.ids.drain());
                last_p.text.push_str(&p.text);
                last_p.tags.append(&mut offset_tags);
                last_p.widgets.append(&mut offset_widgets);
            } else {
                grouped.push(element);
            }
        }

        grouped
    }

    fn group_images(self) -> Element {
        fn push_images(elements: &mut Vec<Element>, images: Vec<Image>) {
            if !images.is_empty() {
                if images.len() == 1 {
                    let mut other = images.into_iter().map(Element::Image).collect::<Vec<_>>();
                    elements.append(&mut other);
                } else {
                    elements.push(Element::Gallery(images));
                }
            }
        }

        fn should_add_to_gallery(image: &Image) -> bool {
            let width = image.width.unwrap_or(200);
            let height = image.height.unwrap_or(200);
            width * height > 35_000
        }

        if let Element::Container(boxed_elements) = self {
            let mut iter = boxed_elements.into_iter().peekable();

            let mut elements = Vec::new();
            let mut images = Vec::new();

            while let Some(element) = iter.next() {
                if let Element::Image(image) = element {
                    if should_add_to_gallery(&image) {
                        if let Some(Element::Image(_next_image)) = iter.peek() {
                            images.push(image);
                            continue;
                        }

                        if !images.is_empty() {
                            images.push(image);
                            push_images(&mut elements, images);
                            images = Vec::new();
                            continue;
                        }
                    }

                    elements.push(Element::Image(image));
                    continue;
                }

                if !images.is_empty() {
                    elements.push(Element::Gallery(images));
                    images = Vec::new();
                }
                elements.push(element);
                continue;
            }

            push_images(&mut elements, images);

            Self::first_or_box(elements).unwrap()
        } else {
            self
        }
    }

    pub fn first_or_box(elements: Vec<Element>) -> Option<Element> {
        let mut iter = elements.into_iter().peekable();

        if let Some(first) = iter.next() {
            if iter.peek().is_none() {
                Some(first)
            } else {
                let elements = std::iter::once(first).chain(iter).collect::<Vec<Element>>();
                Some(Element::Container(elements))
            }
        } else {
            None
        }
    }

    pub fn is_container(&self) -> bool {
        matches!(self, Self::Container(_))
    }
}
