quercus 0.1.1

use std::cmp::max;
use std::fmt;

#[derive(Clone)]
pub struct Tree {
    current_level: usize,
    first_branch: bool,
    text: String,
    style: PipeStyle,
    indent_depth: usize,
    is_pruned: bool,
}

#[derive(Clone)]
pub enum PipeStyle {
    Simple,
    Curvy,
    Thick,
    Double,
}

impl PipeStyle {
    pub fn vertical(&self) -> char {
        match self {
            Self::Simple => '\u{2502}',
            Self::Curvy => '\u{2502}',
            Self::Double => '\u{2551}',
            Self::Thick => '\u{2503}',
        }
    }
    pub fn horizontal(&self) -> char {
        match self {
            Self::Simple => '\u{2500}',
            Self::Curvy => '\u{2500}',
            Self::Double => '\u{2550}',
            Self::Thick => '\u{2501}',
        }
    }
    pub fn t_shape(&self) -> char {
        match self {
            Self::Simple => '\u{251C}',
            Self::Curvy => '\u{251C}',
            Self::Double => '\u{2560}',
            Self::Thick => '\u{2523}',
        }
    }
    pub fn corner(&self) -> char {
        match self {
            Self::Simple => '\u{2514}',
            Self::Curvy => '\u{2570}',
            Self::Double => '\u{255A}',
            Self::Thick => '\u{2517}',
        }
    }
}

impl Default for Tree {
    fn default() -> Self {
        Self {
            current_level: 0,
            first_branch: true,
            text: "".to_string(),
            style: PipeStyle::Simple,
            indent_depth: 2,
            is_pruned: false,
        }
    }
}
impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        <String as fmt::Display>::fmt(&self.text, f)
    }
}

impl Tree {
    pub fn new() -> Self {
        Self::default()
    }

    // ---

    pub fn style(&mut self, style: PipeStyle) -> Self {
        if self.current_level != 0 {
            panic!("Tree can only change style at the root level.")
        }
        self.prune();
        self.style = style;
        self.clone()
    }

    pub fn indent(&mut self, indent: usize) -> Self {
        self.indent_depth = indent;
        self.clone()
    }

    // ---

    pub fn begin_branches(&mut self) {
        if self.first_branch {
            panic!("Cannot begin new branch level without a (root) node.")
        }
        self.current_level += 1;
        self.first_branch = true;
    }

    pub fn node(&mut self, line: &str) {
        self.push_node_line_prefix();

        // Filter-out \n and \r characters
        let filtered_line: String = line.chars().filter(|&c| c != '\n' && c != '\r').collect();

        self.text.push_str(&filtered_line);

        self.first_branch = false;
    }

    pub fn node_many_lines(&mut self, lines: &[&str]) {
        for (idx, line) in lines.iter().enumerate() {
            if idx == 0 {
                self.push_node_line_prefix();
            } else {
                self.push_other_lines_prefix();
            }

            // Filter-out \n and \r characters
            let filtered_line: String = line.chars().filter(|&c| c != '\n' && c != '\r').collect();

            self.text.push_str(&filtered_line);
        }

        self.first_branch = false;
    }

    pub fn end_branches(&mut self) {
        self.current_level -= 1;
        self.first_branch = false;
    }

    pub fn prune(&mut self) -> Self {
        if self.current_level != 0 {
            panic!("Tree can only be pruned at the root level.")
        }

        // Create vector of strings where each string in a line to be printed
        // Lines are in reverse order
        let mut parts: Vec<String> = self.text.split('\n').map(|s| s.to_string()).rev().collect();

        if !self.is_pruned & (parts.len() > 1) {
            // Remove first line if the tree has not yet been pruned.
            parts.pop();
            self.is_pruned = true;
        }

        // Determine the maximum lengh of all lines.
        let mut max_len: usize = 0;

        for s in parts.iter() {
            max_len = max(s.chars().count(), max_len);
        }

        // Append spaces so that all lines have the same length
        for s in parts.iter_mut() {
            s.push_str(" ".to_string().repeat(max_len - s.chars().count()).as_str());
        }

        // Create mutable clone
        let mut copy_parts = parts.clone();

        for (line_nr, line) in parts.iter().enumerate() {
            if line_nr == 0 {
                // Replace T-shapes by corners and vertical pipes by spaces on the last line
                copy_parts[line_nr] = line
                    .chars()
                    .map(|c| {
                        if c == self.style.t_shape() {
                            self.style.corner()
                        } else if c == self.style.vertical() {
                            ' '
                        } else {
                            c
                        }
                    })
                    .collect();
            } else {
                // Replace T-shapes by corners and vertical pipes by spaces if there is no
                // compatible piece on the line below
                copy_parts[line_nr] = line
                    .chars()
                    .zip(copy_parts[line_nr - 1].chars())
                    .map(|(c, r)| {
                        if c == self.style.t_shape()
                            && !(r == self.style.t_shape()
                                || r == self.style.vertical()
                                || r == self.style.corner())
                        {
                            self.style.corner()
                        } else if c == self.style.vertical()
                            && !(r == self.style.t_shape()
                                || r == self.style.vertical()
                                || r == self.style.corner())
                        {
                            ' '
                        } else {
                            c
                        }
                    })
                    .collect();
            }
        }

        // Recover string from vector of lines in the correct order.
        self.text = copy_parts
            .into_iter()
            .rev()
            .collect::<Vec<String>>()
            .join("\n");

        self.clone()
    }

    pub fn prune_and_print(&mut self) {
        println!("{}", self.prune());
    }

    // ---

    fn push_node_line_prefix(&mut self) {
        let mut prefix = String::from("\n");

        for ii in 1..=self.current_level {
            if ii == self.current_level {
                prefix.push_str(format!("{}", self.style.t_shape()).as_str());
                prefix.push_str(
                    self.style
                        .horizontal()
                        .to_string()
                        .repeat(self.indent_depth)
                        .as_str(),
                );
                prefix.push(' ');
            } else {
                prefix.push_str(format!("{}", self.style.vertical()).as_str());
                prefix.push_str(" ".to_string().repeat(self.indent_depth + 1).as_str());
            }
        }

        self.text.push_str(prefix.as_str());
    }

    fn push_other_lines_prefix(&mut self) {
        let mut prefix = String::from("\n");

        for _ in 1..=self.current_level {
            prefix.push_str(format!("{}", self.style.vertical()).as_str());
            prefix.push_str(" ".to_string().repeat(self.indent_depth + 1).as_str());
        }

        self.text.push_str(prefix.as_str());
    }
}